class Udocker(): """Class in charge of managing the udocker binary.""" _CONTAINER_OUTPUT_FILE = SysUtils.join_paths(FileUtils.get_tmp_dir(), "container-stdout") _CONTAINER_NAME = "udocker_container" _SCRIPT_EXEC = "/bin/bash" def __init__(self, lambda_instance): self.lambda_instance = lambda_instance # Create required udocker folder FileUtils.create_folder(SysUtils.get_env_var("UDOCKER_DIR")) # Init the udocker command that will be executed self.udocker_exec = [SysUtils.get_env_var("UDOCKER_EXEC")] self.cont_cmd = self.udocker_exec + ["--quiet", "run"] self.cont_img_id = ConfigUtils.read_cfg_var('container').get('image') if not self.cont_img_id: raise ContainerImageNotFoundError() def _list_udocker_images_cmd(self): return self.udocker_exec + ["images"] def _load_udocker_image_cmd(self): return self.udocker_exec + ["load", "-i", self.cont_img_id] def _download_udocker_image_cmd(self): return self.udocker_exec + ["pull", self.cont_img_id] def _list_udocker_containers_cmd(self): return self.udocker_exec + ["ps"] def _create_udocker_container_cmd(self): return self.udocker_exec + [ "create", f"--name={self._CONTAINER_NAME}", self.cont_img_id ] def _set_udocker_container_execution_mode_cmd(self): return self.udocker_exec + [ "setup", "--execmode=F1", self._CONTAINER_NAME ] def _is_container_image_downloaded(self): cmd_out = SysUtils.execute_cmd_and_return_output( self._list_udocker_images_cmd()) return self.cont_img_id in cmd_out def _load_local_container_image(self): get_logger().info("Loading container image '%s'", self.cont_img_id) SysUtils.execute_cmd(self._load_udocker_image_cmd()) def _download_container_image(self): get_logger().info("Pulling container '%s' from Docker Hub", self.cont_img_id) SysUtils.execute_cmd(self._download_udocker_image_cmd()) def _is_container_available(self): cmd_out = SysUtils.execute_cmd_and_return_output( self._list_udocker_containers_cmd()) return self._CONTAINER_NAME in cmd_out def _create_image(self): if self._is_container_image_downloaded(): get_logger().info("Container image '%s' already available", self.cont_img_id) else: if SysUtils.is_var_in_env("IMAGE_FILE"): self._load_local_container_image() else: self._download_container_image() def _create_container(self): if self._is_container_available(): get_logger().info("Container already available") else: get_logger().info("Creating container based on image '%s'.", self.cont_img_id) SysUtils.execute_cmd(self._create_udocker_container_cmd()) SysUtils.execute_cmd(self._set_udocker_container_execution_mode_cmd()) def _create_command(self): self._add_container_volumes() self._add_container_environment_variables() # Container running script if hasattr(self.lambda_instance, 'script_path'): # Add script in memory as entrypoint self.cont_cmd += [(f"--entrypoint={self._SCRIPT_EXEC} " f"{self.lambda_instance.script_path}"), self._CONTAINER_NAME] # Container with args elif hasattr(self.lambda_instance, 'cmd_args'): # Add args self.cont_cmd += [self._CONTAINER_NAME] self.cont_cmd += self.lambda_instance.cmd_args # Script to be executed every time (if defined) elif hasattr(self.lambda_instance, 'init_script_path'): # Add init script self.cont_cmd += [(f"--entrypoint={self._SCRIPT_EXEC} " f"{self.lambda_instance.init_script_path}"), self._CONTAINER_NAME] # Only container else: self.cont_cmd += [self._CONTAINER_NAME] def _add_container_volumes(self): self.cont_cmd.extend(["-v", SysUtils.get_env_var("TMP_INPUT_DIR")]) self.cont_cmd.extend(["-v", SysUtils.get_env_var("TMP_OUTPUT_DIR")]) self.cont_cmd.extend( ["-v", "/dev", "-v", "/proc", "-v", "/etc/hosts", "--nosysdirs"]) if SysUtils.is_var_in_env('EXTRA_PAYLOAD'): self.cont_cmd.extend(["-v", self.lambda_instance.PERMANENT_FOLDER]) def _add_cont_env_vars(self): for key, value in SysUtils.get_cont_env_vars().items(): self.cont_cmd.extend(_parse_cont_env_var(key, value)) def _add_input_file(self): self.cont_cmd.extend( _parse_cont_env_var("INPUT_FILE_PATH", SysUtils.get_env_var("INPUT_FILE_PATH"))) def _add_output_dir(self): self.cont_cmd.extend( _parse_cont_env_var("TMP_OUTPUT_DIR", SysUtils.get_env_var("TMP_OUTPUT_DIR"))) def _add_storage_object_key(self): self.cont_cmd.extend( _parse_cont_env_var("STORAGE_OBJECT_KEY", SysUtils.get_env_var("STORAGE_OBJECT_KEY"))) def _add_extra_payload_path(self): self.cont_cmd.extend( _parse_cont_env_var("EXTRA_PAYLOAD", SysUtils.get_env_var("EXTRA_PAYLOAD"))) def _add_function_request_id(self): self.cont_cmd.extend( _parse_cont_env_var("REQUEST_ID", self.lambda_instance.get_request_id())) def _add_aws_access_keys(self): self.cont_cmd.extend( _parse_cont_env_var("AWS_ACCESS_KEY_ID", SysUtils.get_env_var("AWS_ACCESS_KEY_ID"))) self.cont_cmd.extend( _parse_cont_env_var("AWS_SECRET_ACCESS_KEY", SysUtils.get_env_var("AWS_SECRET_ACCESS_KEY"))) self.cont_cmd.extend( _parse_cont_env_var("AWS_SESSION_TOKEN", SysUtils.get_env_var("AWS_SESSION_TOKEN"))) def _add_function_ip(self): self.cont_cmd.extend( _parse_cont_env_var("INSTANCE_IP", get_function_ip())) def _add_container_environment_variables(self): self._add_function_request_id() self._add_function_ip() self._add_aws_access_keys() self._add_cont_env_vars() self._add_input_file() self._add_output_dir() self._add_storage_object_key() self._add_extra_payload_path() def prepare_container(self): """Prepares the environment to execute the udocker container.""" self._create_image() self._create_container() self._create_command() def launch_udocker_container(self): """Launches the udocker container. If the execution time of the container exceeds the defined execution time, the container is killed and a warning is raised.""" remaining_seconds = self.lambda_instance.get_remaining_time_in_seconds( ) get_logger().info( "Executing udocker container. Timeout set to '%d' seconds", remaining_seconds) get_logger().debug("Udocker command: '%s'", self.cont_cmd) with open(self._CONTAINER_OUTPUT_FILE, "wb") as out: with subprocess.Popen(self.cont_cmd, stderr=subprocess.STDOUT, stdout=out, start_new_session=True) as process: try: process.wait(timeout=remaining_seconds) except subprocess.TimeoutExpired: get_logger().info("Stopping process '%s'", process) process.kill() raise ContainerTimeoutExpiredWarning() udocker_output = b'' if FileUtils.is_file(self._CONTAINER_OUTPUT_FILE): udocker_output = FileUtils.read_file(self._CONTAINER_OUTPUT_FILE, file_mode="rb") return udocker_output
def test_get_tmp_dir(self, mock_tmp): FileUtils.get_tmp_dir() mock_tmp.assert_called_once()