def bash_script(
script=None,
check_retcode=0,
cwd=None,
env=None,
dbnd_env=True,
output_encoding="utf-8",
popen_kwargs=None,
):
# type:( str, Optional[int],str, Dict[str,str], bool, str, Dict[str,Any]) -> int
# we need a working folder to create bash script
task_run = try_get_current_task_run()
if task_run:
script_dir = str(task_run.local_task_run_root)
else:
script_dir = None
bash_script_path = os.path.join(script_dir, "bash_cmd.sh")
with open(bash_script_path, "wb") as bs:
bs.write(bytes(script, "utf_8"))
log_metric("bash_script", bash_script_path)
logger.info("Bash script location: %s", bash_script_path)
args = ["bash", bash_script_path]
return bash_cmd.func(
args=args,
check_retcode=check_retcode,
cwd=cwd or script_dir,
env=env,
dbnd_env=dbnd_env,
output_encoding=output_encoding,
popen_kwargs=popen_kwargs,
)
def log_operator_result(task_run, result, operator, track_xcom):
_log_result(task_run, result)
# after airflow runs the operator it xcom_push the result, so we log it
if track_xcom and operator.do_xcom_push and result is not None:
from airflow.models import XCOM_RETURN_KEY
log_metric(key=XCOM_RETURN_KEY, value=result)
def t_f_a(t_input, t_param, t_default="d1"):
# type: (DataList[str], str, str) -> DataList[str]
# adds dressing
assert t_default == "d1"
assert t_param == "d2"
log_metric("t_input", len(t_input))
logger.info("Got string: %s", t_input)
return t_input[:2]
def operation_int(input_a,
input_b=0,
pause=0.0,
log_metrics=True,
external_resources=0):
# type: (int, int, float, bool, int) -> int
if log_metrics:
log_metric("input_a", input_a)
log_metric("input_b", input_b)
tr = current_task_run()
for i in range(external_resources):
tr.set_external_resource_urls(
{"url_%s_%d" % (tr.task.task_id, i): "http://localhost"})
if pause:
sleep(pause)
return input_a + input_b
def prepare_docker_for_executor(run, docker_engine):
if docker_engine.container_tag:
logger.info(
"Omitting docker build due to existing container_tag=%s",
docker_engine.container_tag,
)
log_metric("container_tag", docker_engine.container_tag)
return None
config_cls = docker_engine.__class__ # type: Type[ContainerEngineConfig]
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")
if docker_engine.docker_build:
if docker_engine.docker_build_tag:
auto_tag = docker_engine.docker_build_tag
else:
auto_tag = docker_engine.docker_build_tag_base + "_" + CURRENT_TIME_STR
_set_config(config_cls.container_tag, auto_tag)
if (isinstance(docker_engine, KubernetesEngineConfig)
and not docker_engine.docker_build_push):
_set_config(config_cls.image_pull_policy, "Never")
log_metric("docker build tag", auto_tag)
log_metric("container_tag", auto_tag)
docker_build = DockerBuild(
task_name="dbnd_image_build",
image_name=docker_engine.container_repository,
tag=docker_engine.container_tag,
push=docker_engine.docker_build_push,
task_version="now",
task_is_system=True,
)
run.run_dynamic_task(docker_build)
pm.hook.dbnd_build_project_docker(docker_engine=docker_engine,
docker_build_task=docker_build)
return docker_build
else:
logger.info("Omitting docker build due to docker_build=False")
if not docker_engine.container_tag:
raise no_tag_on_no_build()
return None
def write_metrics(a=5):
log_metric("t_f", a)
def t_f_metric(a=5):
log_metric("t_f", a)
def dbnd_sanity_check(check_time=datetime.datetime.now()):
# type: ( datetime.datetime)-> str
logger.info("Running Sanity Check!")
log_metric("Happiness Level", "High")
logger.info("Your system is good to go! Enjoy Databand!")
return "Databand checked at %s" % check_time
def create_model(y_input="y"):
logger.info("Running %s -> operation_y", y_input)
log_metric("a", 1)
return "{} -> operation_y".format(y_input)
def log_metrics(param="param"):
log_metric("param", param)
log_metric("string_metrics", "Pi")
log_metric("float_metric", 3.14159265359)
log_metric("int_metric", 0)
def init_acc_int(input=0):
# type: (int) -> int
log_metric("input", input)
return input
def execute(self, context):
self.log.info("Starting docker container from image %s", self.image)
tls_config = self.__get_tls_config()
if self.docker_conn_id:
self.cli = self.get_hook().get_conn()
else:
self.cli = APIClient(base_url=self.docker_url,
version=self.api_version,
tls=tls_config)
if self.force_pull or len(self.cli.images(name=self.image)) == 0:
for l in self.cli.pull(self.image, stream=True):
try:
output = json.loads(l.decode("utf-8").strip())
if "status" in output:
self.log.info("%s", output["status"])
except Exception:
self.log.info("Failed to parse docker pull status")
# with TemporaryDirectory(prefix='airflowtmp') as host_tmp_dir:
# self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir
# self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir))
self.container = self.cli.create_container(
command=self.get_command(),
environment=self.environment,
host_config=self.cli.create_host_config(
auto_remove=self.auto_remove,
binds=self.volumes,
network_mode=self.network_mode,
shm_size=self.shm_size,
dns=self.dns,
dns_search=self.dns_search,
cpu_shares=int(round(self.cpus * 1024)),
mem_limit=self.mem_limit,
),
image=self.image,
user=self.user,
working_dir=self.working_dir,
)
self.cli.start(self.container["Id"])
log_metric("docker id", self.container["Id"])
line = ""
for line in self.cli.logs(container=self.container["Id"], stream=True):
line = line.strip()
if hasattr(line, "decode"):
line = line.decode("utf-8")
self.log.info(line)
result = self.cli.wait(self.container["Id"])
log_metric("docker status code", result["StatusCode"])
if result["StatusCode"] != 0:
raise AirflowException("docker container failed: " + repr(result))
if self.xcom_push_flag:
return (self.cli.logs(container=self.container["Id"])
if self.xcom_all else str(line))