def emr_operator(self, dag):
spark_submit_command = [
"spark-submit",
"--master",
"yarn",
"--name",
"{{task.task_id}}",
"script.py",
"input.csv",
"output.csv",
]
step_command = dict(
Name="{{task.task_id}}",
ActionOnFailure="CONTINUE",
HadoopJarStep=dict(Jar="command-runner.jar",
Args=spark_submit_command),
)
operator = EmrAddStepsOperator(task_id="emr_task",
job_flow_id=1,
steps=[step_command],
dag=dag)
env = {
"AIRFLOW_CTX_DAG_ID": "test_dag",
"AIRFLOW_CTX_EXECUTION_DATE": "emr_task",
"AIRFLOW_CTX_TASK_ID": "1970-01-01T0000.000",
"AIRFLOW_CTX_TRY_NUMBER": "1",
"AIRFLOW_CTX_UID": get_airflow_instance_uid(),
}
with wrap_operator_with_tracking_info(env, operator):
return operator
def extract_airflow_context(airflow_context):
# type: (Dict[str, Any]) -> Optional[AirflowTaskContext]
"""Create AirflowTaskContext for airflow_context dict"""
task_instance = airflow_context.get("task_instance")
if task_instance is None:
return None
dag_id = task_instance.dag_id
task_id = task_instance.task_id
execution_date = str(task_instance.execution_date)
try_number = task_instance.try_number
if dag_id and task_id and execution_date:
return AirflowTaskContext(
dag_id=dag_id,
execution_date=execution_date,
task_id=task_id,
try_number=try_number,
context=airflow_context,
airflow_instance_uid=get_airflow_instance_uid(),
)
logger.debug(
"airflow context from inspect, at least one of those params is missing"
"dag_id: {}, execution_date: {}, task_id: {}".format(
dag_id, execution_date, task_id))
return None
def get_plugin_metadata(self) -> PluginMetadata:
from airflow import version as airflow_version
import dbnd_airflow
from dbnd_airflow.export_plugin.compat import get_api_mode
return PluginMetadata(
airflow_version=airflow_version.version,
plugin_version=dbnd_airflow.__version__ + " v2",
airflow_instance_uid=get_airflow_instance_uid(),
api_mode=get_api_mode(),
)
def get_meta(metrics):
import flask
meta = AirflowExportMeta(
airflow_version=airflow_version,
plugin_version=" ".join([dbnd_airflow.__version__, "v2"]),
airflow_instance_uid=get_airflow_instance_uid(),
api_mode=get_api_mode(),
request_args=dict(flask.request.args)
if flask.has_request_context() else {},
metrics={
"performance": metrics.get("perf_metrics", {}),
"sizes": metrics.get("size_metrics", {}),
},
)
return meta
def context_to_airflow_vars(context, in_env_var_format=False):
# original_context_to_airflow_vars is created during function override in patch_models()
params = airflow.utils.operator_helpers._original_context_to_airflow_vars(
context=context, in_env_var_format=in_env_var_format)
if in_env_var_format:
task_run = try_get_current_task_run() # type: TaskRun
if task_run:
params = extend_airflow_ctx_with_dbnd_tracking_info(
task_run, params)
try_number = str(context['task_instance'].try_number)
params.update({
"AIRFLOW_CTX_TRY_NUMBER": try_number,
"AIRFLOW_CTX_UID": get_airflow_instance_uid()
})
return params
def spark_submit_operator(self, dag):
operator = SparkSubmitOperator(
task_id="spark_submit_task",
application="script.py",
application_args=["input.csv", "output.csv"],
dag=dag,
)
env = {
"AIRFLOW_CTX_DAG_ID": "test_dag",
"AIRFLOW_CTX_EXECUTION_DATE": "spark_submit_task",
"AIRFLOW_CTX_TASK_ID": "1970-01-01T0000.000",
"AIRFLOW_CTX_TRY_NUMBER": "1",
"AIRFLOW_CTX_UID": get_airflow_instance_uid(),
}
with wrap_operator_with_tracking_info(env, operator):
return operator
def extract_airflow_conf(context):
task_instance = context.get("task_instance")
if task_instance is None:
return {}
dag_id = task_instance.dag_id
task_id = task_instance.task_id
execution_date = str(task_instance.execution_date)
try_number = str(task_instance.try_number)
if dag_id and task_id and execution_date:
return {
"AIRFLOW_CTX_DAG_ID": dag_id,
"AIRFLOW_CTX_EXECUTION_DATE": execution_date,
"AIRFLOW_CTX_TASK_ID": task_id,
"AIRFLOW_CTX_TRY_NUMBER": try_number,
"AIRFLOW_CTX_UID": get_airflow_instance_uid(),
}
return {}
def get_airflow_conf(
dag_id="{{dag.dag_id}}",
task_id="{{task.task_id}}",
execution_date="{{ts}}",
try_number="{{task_instance._try_number}}",
):
"""
These properties are
AIRFLOW_CTX_DAG_ID - name of the Airflow DAG to associate a run with
AIRFLOW_CTX_EXECUTION_DATE - execution_date to associate a run with
AIRFLOW_CTX_TASK_ID - name of the Airflow Task to associate a run with
AIRFLOW_CTX_TRY_NUMBER - try number of the Airflow Task to associate a run with
"""
airflow_conf = {
"AIRFLOW_CTX_DAG_ID": dag_id,
"AIRFLOW_CTX_EXECUTION_DATE": execution_date,
"AIRFLOW_CTX_TASK_ID": task_id,
"AIRFLOW_CTX_TRY_NUMBER": try_number,
"AIRFLOW_CTX_UID": get_airflow_instance_uid(),
}
airflow_conf.update(get_databand_url_conf())
return airflow_conf
def execute(self, context):
logger.debug("Running dbnd dbnd_task from airflow operator %s",
self.task_id)
dag = context["dag"]
execution_date = context["execution_date"]
dag_id = dag.dag_id
run_uid = get_job_run_uid(
airflow_instance_uid=get_airflow_instance_uid(),
dag_id=dag_id,
execution_date=execution_date,
)
# Airflow has updated all relevant fields in Operator definition with XCom values
# now we can create a real dbnd dbnd_task with real references to dbnd_task
new_kwargs = {}
for p_name in self.dbnd_task_params_fields:
new_kwargs[p_name] = getattr(self, p_name, None)
# this is the real input value after
if p_name in self.dbnd_xcom_inputs:
new_kwargs[p_name] = target(new_kwargs[p_name])
new_kwargs["_dbnd_disable_airflow_inplace"] = True
dag_ctrl = self.get_dbnd_dag_ctrl()
with DatabandContext.context(_context=dag_ctrl.dbnd_context) as dc:
logger.debug("Running %s with kwargs=%s ", self.task_id,
new_kwargs)
dbnd_task = dc.task_instance_cache.get_task_by_id(
self.dbnd_task_id)
# rebuild task with new values coming from xcom->operator
with dbnd_task.ctrl.task_context(phase=TaskContextPhase.BUILD):
dbnd_task = dbnd_task.clone(
output_params_to_clone=self.dbnd_overridden_output_params,
**new_kwargs)
logger.debug("Creating inplace databand run for driver dump")
dag_task = Task(task_name=dag.dag_id,
task_target_date=execution_date)
dag_task.set_upstream(dbnd_task)
# create databand run
with new_databand_run(
context=dc,
job_name=dag.dag_id,
run_uid=run_uid,
existing_run=False,
) as dr: # type: DatabandRun
dr.run_executor = run_executor = RunExecutor(
run=dr,
root_task_or_task_name=dag_task,
send_heartbeat=False)
run_executor._init_task_runs_for_execution(
task_engine=run_executor.local_engine)
# dr.driver_task_run.set_task_run_state(state=TaskRunState.RUNNING)
# "make dag run"
# dr.root_task_run.set_task_run_state(state=TaskRunState.RUNNING)
dbnd_task_run = dr.get_task_run_by_id(dbnd_task.task_id)
needs_databand_run_save = dbnd_task._conf__require_run_dump_file
if needs_databand_run_save:
run_executor.save_run_pickle()
logger.info(
dbnd_task.ctrl.banner("Running task '%s'." %
dbnd_task.task_name,
color="cyan"))
# should be replaced with tr._execute call
dbnd_task_run.runner.execute(airflow_context=context,
handle_sigterm=False)
logger.debug("Finished to run %s", self)
result = {
output_name:
convert_to_safe_types(getattr(dbnd_task, output_name))
for output_name in self.dbnd_xcom_outputs
}
return result