def test_dataloadingconfig(): dlc = task.DataLoadingConfig("s3://input/path", "s3://output/path", True, task.DataLoadingConfig.LITERALMAP_FORMAT_YAML) dlc2 = task.DataLoadingConfig.from_flyte_idl(dlc.to_flyte_idl()) assert dlc2 == dlc dlc = task.DataLoadingConfig("s3://input/path", "s3://output/path", True, task.DataLoadingConfig.LITERALMAP_FORMAT_YAML, io_strategy=task.IOStrategy()) dlc2 = task.DataLoadingConfig.from_flyte_idl(dlc.to_flyte_idl()) assert dlc2 == dlc
def get_container( self, settings: SerializationSettings) -> _task_model.Container: env = { **settings.env, **self.environment } if self.environment else settings.env return _get_container_definition( image=self._image, command=self._cmd, args=self._args, data_loading_config=_task_model.DataLoadingConfig( input_path=self._input_data_dir, output_path=self._output_data_dir, format=self._md_format.value, enabled=True, io_strategy=self._io_strategy.value if self._io_strategy else None, ), environment=env, cpu_request=self.resources.requests.cpu, cpu_limit=self.resources.limits.cpu, memory_request=self.resources.requests.mem, memory_limit=self.resources.limits.mem, ephemeral_storage_request=self.resources.requests. ephemeral_storage, ephemeral_storage_limit=self.resources.limits.ephemeral_storage, )
def get_container(self, settings: SerializationSettings) -> _task_model.Container: env = {**settings.env, **self.environment} if self.environment else settings.env return _get_container_definition( image=self._image, command=self._cmd, args=self._args, data_loading_config=_task_model.DataLoadingConfig( input_path=self._input_data_dir, output_path=self._output_data_dir, format=self._md_format.value, enabled=True, io_strategy=self._io_strategy.value if self._io_strategy else None, ), environment=env, )
def __init__( self, inputs: Dict[str, FlyteSdkType], image: str, outputs: Dict[str, FlyteSdkType] = None, input_data_dir: str = None, output_data_dir: str = None, metadata_format: int = METADATA_FORMAT_JSON, io_strategy: _task_models.IOStrategy = None, command: List[str] = None, args: List[str] = None, storage_request: str = None, cpu_request: str = None, gpu_request: str = None, memory_request: str = None, storage_limit: str = None, cpu_limit: str = None, gpu_limit: str = None, memory_limit: str = None, environment: Dict[str, str] = None, interruptible: bool = False, discoverable: bool = False, discovery_version: str = None, retries: int = 1, timeout: _datetime.timedelta = None, ): """ :param inputs: :param outputs: :param image: :param command: :param args: :param storage_request: :param cpu_request: :param gpu_request: :param memory_request: :param storage_limit: :param cpu_limit: :param gpu_limit: :param memory_limit: :param environment: :param interruptible: :param discoverable: :param discovery_version: :param retries: :param timeout: :param input_data_dir: This is the directory where data will be downloaded to :param output_data_dir: This is the directory where data will be uploaded from :param metadata_format: Format in which the metadata will be available for the script """ # Set as class fields which are used down below to configure implicit # parameters self._data_loading_config = _task_models.DataLoadingConfig( input_path=input_data_dir, output_path=output_data_dir, format=metadata_format, enabled=True, io_strategy=io_strategy, ) metadata = _task_models.TaskMetadata( discoverable, # This needs to have the proper version reflected in it _task_models.RuntimeMetadata( _task_models.RuntimeMetadata.RuntimeType.FLYTE_SDK, flytekit.__version__, "python", ), timeout or _datetime.timedelta(seconds=0), _literals.RetryStrategy(retries), interruptible, discovery_version, None, ) # The interface is defined using the inputs and outputs i = _interface.TypedInterface(inputs=types_to_variable(inputs), outputs=types_to_variable(outputs)) # This sets the base SDKTask with container etc super(SdkRawContainerTask, self).__init__( _constants.SdkTaskType.RAW_CONTAINER_TASK, metadata, i, None, container=_get_container_definition( image=image, args=args, command=command, data_loading_config=self._data_loading_config, storage_request=storage_request, cpu_request=cpu_request, gpu_request=gpu_request, memory_request=memory_request, storage_limit=storage_limit, cpu_limit=cpu_limit, gpu_limit=gpu_limit, memory_limit=memory_limit, environment=environment, ), )