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"