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)