Add kubernetes template for airflow + more

Dockerfile_k8s

@@ -2,7 +2,7 @@ FROM docker.io/bitnami/spark:3.5.0
 USER root
 
 RUN python -m pip install --upgrade pip
-RUN pip3 install --no-deps yaetos==0.12.0
+RUN pip3 install --no-deps yaetos==0.12.1
 # Force latest version to avoid using previous ones.
 RUN pip3 install -r /opt/bitnami/python/lib/python3.11/site-packages/yaetos/scripts/requirements_base.txt
 # Installing libraries required by Yaetos and more. Using this since requirements_base.txt has exact versions.

conf/jobs_metadata.yml

@@ -378,6 +378,16 @@ jobs:
     output: {'path':'{{base_path}}/wiki_example/output_ex20_filelist/{{now}}/dataset.csv', 'type':'csv', 'df_type':'pandas'}
     spark_boot: False
 
+  examples/ex21_deploy_airflow_job:
+    description: "Test deployment to airflow."
+    py_job: jobs/generic/copy_job.py
+    inputs:
+      table_to_copy: {'path':"tests/fixtures/data_sample/wiki_example/input/dataset.csv", 'type':'csv'}
+    output: {'path':'{{base_path}}/load_example/test_files/{{now}}/dataset.csv', 'type':'csv'}
+    s3_dags: 's3://mylake-dev/pipelines_metadata/airflow_dags'
+    local_dags: './airflow_dags/'
+    spark_boot: False
+
   # wordcount_raw_job: #Job exists but doesn't rely on jobs_metadata entries
 
   # ----- Marketing Jobs --------
@@ -458,6 +468,16 @@ common_params:
     redshift_s3_tmp_dir: s3a://dev-spark/tmp_spark/
     email_cred_section: some_email_cred_section  # Section from "connection_file"
     spark_version: '3.5' # options: '2.4', '3.0', '3.4' or '3.5'
+    k8s_url: 'k8s://https://kubernetes.docker.internal:6443'
+    k8s_name: 'my-pyspark-job'
+    k8s_executor_instances: '2'
+    k8s_namespace: 'a_k8s_namespace'
+    k8s_image_service: 'a_k8s_image_service'
+    k8s_upload_path: 's3a://a_k8s_upload_path'
+    k8s_driver_podTemplateFile: conf/k8s_setup_spark_submit_driver.yaml
+    k8s_executor_podTemplateFile: conf/k8s_setup_spark_submit_executor.yaml
+    aws_region: eu-west-1  # TODO: remove this.
+    k8s_podname: a_podname  # TODO: make it nullable so jobs can rerun back to back.
     default_aws_modes: 'dev_EMR'
     default_local_modes: 'dev_local'
     aws_modes: ['dev_EMR','prod_EMR']

jobs/examples/ex15_copy_job_multi_path.py

@@ -43,7 +43,7 @@ def expand_input_path(self, path, **kwargs):
         base_path = self.jargs.base_path
         path_partly_expanded = path.replace('{category}', category) \
                                    .replace('{subcategory}', subcategory)
-        path = Path_Handler(path_partly_expanded, base_path, self.jargs.merged_args.get('root_path')).expand_later()
+        path = Path_Handler(path_partly_expanded, base_path, self.jargs.merged_args.get('root_path')).expand_latest()
         return path
 
     def expand_output_path(self, path, now_dt, **kwargs):

jobs/generic/copy_raw_job.py

@@ -1,7 +1,7 @@
 """
 Job meant to run locally to get data from AWS S3 to local. Updates required to run in cluster.
 """
-from yaetos.etl_utils import ETL_Base, Commandliner, get_aws_setup
+from yaetos.etl_utils import ETL_Base, Commandliner, get_aws_setup  # FS_Ops_Dispatcher
 import os
 from cloudpathlib import CloudPath as CPt
 import fnmatch
@@ -27,14 +27,17 @@ def transform(self, files_to_copy):
             pattern = '*'
             pattern_type = 'glob'
 
+        # TODO: replace code below (and all functions) with the commented code
+        # FS_Ops_Dispatcher().copy_file(path_in, path_out)
+
         session = get_aws_setup(self.jargs.merged_args)
         s3 = session.client('s3')
 
         file_number = self.get_size(s3, path_raw_in.bucket, path_raw_in.key, pattern, pattern_type)
         self.logger.info(f"Number of files to be downloaded {file_number}")
 
         self.download_files(s3, path_raw_in.bucket, path_raw_in.key, pattern, pattern_type, path_raw_out)
-        self.logger.info("Finished downloading all files")
+        self.logger.info("Finished copying all files")
         return None
 
     def download_files(self, s3, bucket_name, prefix, pattern, pattern_type, path_raw_out):

jobs/generic/list_files_job.py

@@ -1,4 +1,4 @@
-from yaetos.etl_utils import ETL_Base, Commandliner, get_aws_setup
+from yaetos.etl_utils import ETL_Base, Commandliner, get_aws_setup  # FS_Ops_Dispatcher
 from cloudpathlib import CloudPath as CPt
 import fnmatch
 import re
@@ -26,6 +26,9 @@ def transform(self, files):
             pattern = '*'
             pattern_type = 'glob'
 
+        # TODO: replace code below (and all functions) with the commented code
+        # files = FS_Ops_Dispatcher().list_files(path, regex=None, globy=None)
+
         session = get_aws_setup(self.jargs.merged_args)
 
         s3 = session.client('s3')

tests/fixtures/ref_airflow_emr_job_dag.py

@@ -0,0 +1,140 @@
+
+from airflow import DAG
+from airflow.providers.amazon.aws.operators.emr import EmrCreateJobFlowOperator, EmrAddStepsOperator  # EmrTerminateJobFlowOperator
+from airflow.providers.amazon.aws.sensors.emr import EmrStepSensor  # EmrJobFlowSensor
+from airflow.utils.dates import days_ago  # noqa: F401
+from datetime import timedelta
+import dateutil
+
+
+DAG_ARGS = {
+    'dag_id': 'ex-job_x',
+    'dagrun_timeout': timedelta(hours=2),
+    'start_date': dateutil.parser.parse("2024-07-15T00:00:00+00:00"),  # ignore_in_diff
+    'schedule': '@once',
+    'tags': ['emr'],
+    'default_args': {
+        'owner': 'me',
+        'depends_on_past': False,
+        'email': [],
+        'email_on_failure': False,
+        'email_on_retry': False,
+    },
+}
+
+
+CLUSTER_JOB_FLOW_OVERRIDES = {
+    'Name': 'yaetos__ex_s_job_x__20240101T000000',  # ignore_in_diff
+    'ReleaseLabel': 'emr-6.1.1',
+    'Applications': [{'Name': 'Hadoop'}, {'Name': 'Spark'}],
+    'Instances': {
+        'InstanceGroups': [
+            {
+                'Name': "Main nodes",
+                'Market': 'ON_DEMAND',
+                'InstanceRole': 'MASTER',
+                'InstanceType': 'm5.xlarge',
+                'InstanceCount': 1,
+            },
+
+            {
+                'Name': 'Secondary nodes',
+                'Market': 'ON_DEMAND',
+                'InstanceRole': 'CORE',
+                'InstanceType': 'm5.xlarge',
+                'InstanceCount': '2',
+            }
+
+        ],
+        'KeepJobFlowAliveWhenNoSteps': False,
+        'TerminationProtected': False,
+        'Ec2KeyName': 'to_be_filled',
+        'Ec2SubnetId': 'to_be_filled',
+        # 'AdditionalMasterSecurityGroups': extra_security_gp,  # TODO : make optional in future. "[self.extra_security_gp] if self.extra_security_gp else []" doesn't work.
+    },
+    'VisibleToAllUsers': True,
+    'JobFlowRole': 'EMR_EC2_DefaultRole',
+    'ServiceRole': 'EMR_DefaultRole',
+    'LogUri': "s3://mylake-dev/pipelines_metadata/manual_run_logs/",
+    'BootstrapActions': [{
+        'Name': 'setup_nodes',
+        'ScriptBootstrapAction': {
+            'Path': 's3n://mylake-dev/pipelines_metadata/jobs_code/yaetos__ex_s_job_x__20240701T000000/code_package/setup_nodes.sh',  # ignore_in_diff
+            'Args': []
+        }
+    }],
+    'Configurations': [
+        {  # Section to force python3 since emr-5.x uses python2 by default.
+            "Classification": "spark-env",
+            "Configurations": [{
+                "Classification": "export",
+                "Properties": {"PYSPARK_PYTHON": "/usr/bin/python3"}
+            }]
+        },
+        # { # Section to add jars (redshift...), not used for now, since passed in spark-submit args.
+        # "Classification": "spark-defaults",
+        # "Properties": { "spark.jars": ["/home/hadoop/redshift_tbd.jar"], "spark.driver.memory": "40G", "maximizeResourceAllocation": "true"},
+        # }
+    ]
+}
+
+EMR_STEPS = [
+    {
+        'Name': 'Run Setup',
+        'ActionOnFailure': 'CONTINUE',
+        'HadoopJarStep': {
+            'Jar': 's3://to_be_filled.elasticmapreduce/libs/script-runner/script-runner.jar',
+            'Args': [
+                "s3://mylake-dev/pipelines_metadata/jobs_code/yaetos__ex_s_job_x__20240701T000000/code_package/setup_master.sh",  # ignore_in_diff
+                "s3://mylake-dev/pipelines_metadata/jobs_code/yaetos__ex_s_job_x__20240701T000000/code_package",  # ignore_in_diff
+            ]
+        }
+    },
+    {
+        'Name': 'Spark Application',
+        'ActionOnFailure': 'CONTINUE',
+        'HadoopJarStep': {
+            'Jar': 'command-runner.jar',
+            'Args': ['spark-submit', '--verbose', '--py-files=/home/hadoop/app/scripts.zip', '/home/hadoop/app/some/job.py', '--mode=None', '--deploy=none', '--storage=s3', '--job_name=ex/job_x'],
+        },
+    }
+]
+
+with DAG(**DAG_ARGS) as dag:
+
+    cluster_creator = EmrCreateJobFlowOperator(
+        task_id='start_emr_cluster',
+        aws_conn_id='aws_default',
+        emr_conn_id='emr_default',
+        job_flow_overrides=CLUSTER_JOB_FLOW_OVERRIDES
+    )
+
+    step_adder = EmrAddStepsOperator(
+        task_id='add_steps',
+        job_flow_id="{{ task_instance.xcom_pull(task_ids='start_emr_cluster', key='return_value') }}",
+        aws_conn_id='aws_default',
+        steps=EMR_STEPS,
+    )
+
+    step_checker = EmrStepSensor(
+        task_id='watch_step',
+        job_flow_id="{{ task_instance.xcom_pull('start_emr_cluster', key='return_value') }}",
+        step_id="{{ task_instance.xcom_pull(task_ids='add_steps', key='return_value')[1] }}",  # [1] to watch 2nd step, the spark application.
+        aws_conn_id='aws_default',
+    )
+
+    # # not used for now
+    # cluster_checker = EmrJobFlowSensor(
+    #     task_id='check_cluster',
+    #     job_flow_id="{{ task_instance.xcom_pull('start_emr_cluster', key='return_value') }}",
+    #     aws_conn_id='aws_default',
+    # )
+
+    # # not used for now
+    # terminate_cluster = EmrTerminateJobFlowOperator(
+    #     task_id='terminate_cluster',
+    #     job_flow_id="{{ task_instance.xcom_pull('start_emr_cluster', key='return_value') }}",
+    #     aws_conn_id='aws_default',
+    # )
+
+    cluster_creator >> step_adder >> step_checker

tests/fixtures/ref_airflow_k8s_job_dag.py

@@ -0,0 +1,45 @@
+
+from airflow import DAG
+from airflow.providers.cncf.kubernetes.operators.spark_kubernetes import SparkKubernetesOperator
+from airflow.providers.cncf.kubernetes.sensors.spark_kubernetes import SparkKubernetesSensor
+from airflow.utils.dates import days_ago  # noqa: F401
+from datetime import timedelta
+import dateutil
+
+
+DAG_ARGS = {
+    'dag_id': 'ex-job_x',
+    'dagrun_timeout': timedelta(hours=2),
+    'start_date': dateutil.parser.parse("2024-07-15T00:00:00+00:00"),  # ignore_in_diff
+    'schedule': '@once',
+    'tags': ['emr'],
+    'default_args': {
+        'owner': 'me',
+        'depends_on_past': False,
+        'email': [],
+        'email_on_failure': False,
+        'email_on_retry': False,
+    },
+}
+
+
+with DAG(**DAG_ARGS) as dag:
+
+    spark_submit = SparkKubernetesOperator(
+        task_id='spark_submit_task',
+        namespace='None',
+        application_file='None',
+        kubernetes_conn_id='k8s_default',
+        do_xcom_push=True,
+    )
+
+    spark_sensor = SparkKubernetesSensor(
+        task_id='watch_step',
+        namespace='None',
+        application_name="{{ task_instance.xcom_pull(task_ids='spark_submit_task')['metadata']['name'] }}",
+        kubernetes_conn_id='k8s_default',
+        poke_interval=60,
+        timeout=600,
+    )
+
+    spark_submit >> spark_sensor