def execute_callable(self):
with TemporaryDirectory(prefix='venv') as tmp_dir:
if self.templates_dict:
self.op_kwargs['templates_dict'] = self.templates_dict
input_filename = os.path.join(tmp_dir, 'script.in')
output_filename = os.path.join(tmp_dir, 'script.out')
string_args_filename = os.path.join(tmp_dir, 'string_args.txt')
script_filename = os.path.join(tmp_dir, 'script.py')
prepare_virtualenv(venv_directory=tmp_dir,
python_bin=f'python{self.python_version}'
if self.python_version else None,
system_site_packages=self.system_site_packages,
requirements=self.requirements)
self._write_args(input_filename)
self._write_string_args(string_args_filename)
write_python_script(jinja_context=dict(
op_args=self.op_args,
op_kwargs=self.op_kwargs,
pickling_library=self.pickling_library.__name__,
python_callable=self.python_callable.__name__,
python_callable_source=dedent(
inspect.getsource(self.python_callable))),
filename=script_filename)
execute_in_subprocess(cmd=[
f'{tmp_dir}/bin/python', script_filename, input_filename,
output_filename, string_args_filename
])
return self._read_result(output_filename)
def execute(self, context: 'Context'):
with TemporaryDirectory(prefix='venv') as tmp_dir:
input_filename = os.path.join(tmp_dir, 'script.in')
script_filename = os.path.join(tmp_dir, 'script.py')
with open(input_filename, 'wb') as file:
if self.op_args or self.op_kwargs:
self.pickling_library.dump({'args': self.op_args, 'kwargs': self.op_kwargs}, file)
py_source = self._get_python_source()
write_python_script(
jinja_context=dict(
op_args=self.op_args,
op_kwargs=self.op_kwargs,
pickling_library=self.pickling_library.__name__,
python_callable=self.python_callable.__name__,
python_callable_source=py_source,
string_args_global=False,
),
filename=script_filename,
)
# Pass the python script to be executed, and the input args, via environment variables. This is
# more than slightly hacky, but it means it can work when Airflow itself is in the same Docker
# engine where this task is going to run (unlike say trying to mount a file in)
self.environment["__PYTHON_SCRIPT"] = _b64_encode_file(script_filename)
if self.op_args or self.op_kwargs:
self.environment["__PYTHON_INPUT"] = _b64_encode_file(input_filename)
else:
self.environment["__PYTHON_INPUT"] = ""
self.command = self.generate_command()
return super().execute(context)
def execute_callable(self):
with TemporaryDirectory(prefix='venv') as tmp_dir:
requirements_file_name = f'{tmp_dir}/requirements.txt'
if not isinstance(self.requirements, str):
requirements_file_contents = "\n".join(
str(dependency) for dependency in self.requirements)
else:
requirements_file_contents = self.requirements
if not self.system_site_packages and self.use_dill:
requirements_file_contents += '\ndill'
with open(requirements_file_name, 'w') as file:
file.write(requirements_file_contents)
if self.templates_dict:
self.op_kwargs['templates_dict'] = self.templates_dict
input_filename = os.path.join(tmp_dir, 'script.in')
output_filename = os.path.join(tmp_dir, 'script.out')
string_args_filename = os.path.join(tmp_dir, 'string_args.txt')
script_filename = os.path.join(tmp_dir, 'script.py')
prepare_virtualenv(
venv_directory=tmp_dir,
python_bin=f'python{self.python_version}'
if self.python_version else None,
system_site_packages=self.system_site_packages,
requirements_file_path=requirements_file_name,
pip_install_options=self.pip_install_options,
)
self._write_args(input_filename)
self._write_string_args(string_args_filename)
write_python_script(
jinja_context=dict(
op_args=self.op_args,
op_kwargs=self.op_kwargs,
pickling_library=self.pickling_library.__name__,
python_callable=self.python_callable.__name__,
python_callable_source=self.get_python_source(),
),
filename=script_filename,
render_template_as_native_obj=self.dag.
render_template_as_native_obj,
)
execute_in_subprocess(cmd=[
f'{tmp_dir}/bin/python',
script_filename,
input_filename,
output_filename,
string_args_filename,
])
return self._read_result(output_filename)
