def execute(self, context):
scheduled_run_info = ScheduledRunInfo(
scheduled_job_uid=self.scheduled_job_uid,
scheduled_job_dag_run_id=context.get("dag_run").id,
scheduled_date=context.get("task_instance").execution_date,
)
# disable heartbeat at this level,
# otherwise scheduled jobs that will run on Kubernetes
# will always have a heartbeat even if the actual driver
# sent to Kubernetes is lost down the line,
# which is the main purpose of the heartbeat
with new_dbnd_context(
name="airflow",
conf={
RunConfig.task_executor_type:
override(TaskExecutorType.local),
RunConfig.parallel: override(False),
LoggingConfig.disabled: override(True),
},
) as dc:
launcher_task = Launcher(
scheduled_cmd=self.scheduled_cmd,
task_name=context.get("dag").dag_id,
task_version="now",
task_is_system=True,
shell=self.shell,
)
dc.dbnd_run_task(
task_or_task_name=launcher_task,
scheduled_run_info=scheduled_run_info,
send_heartbeat=False,
)
def _save_graph(self, task):
with new_dbnd_context(
conf={
RunConfig.task_executor_type: override(
TaskExecutorType.local),
CoreConfig.tracker: override(["console"]),
}) as dc:
run = dc.dbnd_run_task(task_or_task_name=task)
run.save_run()
loaded_run = DatabandRun.load_run(dump_file=run.driver_dump,
disable_tracking_api=False)
assert loaded_run
return run
def test_with_task_class_extend_and_override_layer(self, config_sections):
with config(config_sections):
with config({DummyConfig.dict_config: override({"override": "win"})}):
# can't merge into override value
example_task_with_task_config_extend.task(
expected_dict={"override": "win"}
).dbnd_run()
def launch(self, context, scheduled_run_info: ScheduledRunInfo):
with new_dbnd_context(
name="airflow",
conf={
RunConfig.task_executor_type:
override(TaskExecutorType.local),
RunConfig.parallel: override(False),
LoggingConfig.disabled: override(True),
},
) as dc:
launcher_task = Launcher(
scheduled_cmd=self.scheduled_cmd,
task_name=context.get("dag").dag_id,
task_version="now",
task_is_system=True,
shell=self.shell,
)
dc.dbnd_run_task(
task_or_task_name=launcher_task,
scheduled_run_info=scheduled_run_info,
send_heartbeat=False,
)
from dbnd._core.configuration.config_readers import read_from_config_stream
from dbnd._core.run.databand_run import DatabandRun
from dbnd._core.settings import CoreConfig
from dbnd._core.task_build.task_context import TaskContextPhase
from dbnd._core.task_build.task_registry import build_task_from_config
from dbnd._core.task_run.task_run import TaskRun
from dbnd._core.utils import seven
def dynamically_calculated():
logging.error("Calculating config value for test! test_f_value")
return "test_f_value"
@task(task_config=dict(kubernetes=dict(
cluster_context=override("test"),
limits={"test_limits": 1},
container_tag=
"@python://test_dbnd.task.task_configuration.test_task_config.dynamically_calculated",
)))
def dummy_nested_config_task(expected, config_name):
# type: ( object, str)-> object
# explicitly build config for k8s
actual = build_task_from_config(task_name=config_name)
return (actual.limits, actual.cluster_context, actual.container_tag)
@task(task_config=dict(kubernetes=dict(
labels=extend({"task_name": "the one and only"}))))
from dbnd import override, task
from dbnd._core.run.databand_run import DatabandRun
from dbnd._core.task_build.task_context import TaskContextPhase
from dbnd._core.task_build.task_registry import build_task_from_config
from dbnd._core.task_run.task_run import TaskRun
from dbnd._core.utils.seven import qualname_func
def dynamically_calculated():
logging.error("Calculating!")
return "test_f_value"
@task(task_config=dict(tconfig=dict(
config_value_s1=override("override_config_s1"),
config_value_s2="regular_config_s2",
)))
def dummy_nested_config_task(config_name):
# type: ( str)-> object
actual = build_task_from_config(task_name=config_name)
return (actual.config_value_s1, actual.config_value_s2)
class TestTaskConfig(object):
def test_task_config_override(self):
actual = dummy_nested_config_task.dbnd_run(config_name="tconfig")
assert actual.task.result.load(object) == (
"override_config_s1",
"regular_config_s2",
)
def _set_config(parameter, value):
# set value in already existing object
setattr(docker_engine, parameter.name, value)
dbnd_config.set_parameter(parameter,
override(value),
source="prepare_docker_for_executor")
# _type = dummy
# dict_config = {"second_key": "env_2_value"}
# """
@task
def example_task(name, expected_dict):
assert Config.from_databand_context(name).dict_config == expected_dict
@task(task_config={"dummy": {"dict_config": {"A": "a"}}})
def example_task_with_task_config(expected_dict):
assert DummyConfig.from_databand_context().dict_config == expected_dict
@task(task_config={"dummy": {"dict_config": override({"A": "a"})}})
def example_task_with_task_config_override(expected_dict):
assert DummyConfig.from_databand_context().dict_config == expected_dict
@task(task_config={"dummy": {"dict_config": extend({"new_key": "value"})}})
def example_task_with_task_config_extend(expected_dict):
assert DummyConfig.from_databand_context().dict_config == expected_dict
class TestBuildTaskWithMergeOperator(object):
@pytest.fixture
def config_sections(self):
return {
DummyConfig.dict_config: {"first_key": "basic_value"},
# DummyConfig.list_config: ["basic_value"],
import datetime
import logging
from time import sleep
from dbnd import config, databand_lib_path, override, pipeline, task
from dbnd.tasks.basics import dbnd_sanity_check
from dbnd_docker.docker.docker_task import DockerRunTask
logger = logging.getLogger(__name__)
@task(task_config=dict(kubernetes=dict(
pod_yaml=override(
databand_lib_path("conf", "kubernetes-pod-tensorflow.yaml")),
trap_exit_file_flag=override("/output/training_logs/main-terminated"),
limits=override({"nvidia.com/gpu": 1}),
)))
def task_with_custom_k8s_yml_gpu(
check_time=datetime.datetime.now(), sleep_time_sec=120):
# type: ( datetime.datetime, int)-> str
config.log_current_config(as_table=True)
logger.info("Running Kube Sanity Check!")
if sleep_time_sec:
logger.info("sleeping for %s", sleep_time_sec)
sleep(sleep_time_sec)
return "Databand checked at %s" % check_time
class ExampleDockerNativeTask(DockerRunTask):
command = "echo hi"