Raw Container Task Local Execution

[Future-Outlier]

Tracking issue

flyteorg/flyte#3876

Why are the changes needed?

We want to support users to run raw-container task locally.

Note: I didn't support datetime.timedelta, List and Dict.
Since I found that copilot has bug to deal with that, so I will support them until those bugs are fixed.
(This implies that we have to fix flytecopilot for nested types and collections types first)

What changes were proposed in this pull request?

1. override local_execute function in class ContainerTask(PythonTask)
2. use docker python library to execute docker command
3. support int, str, float, bool, datetime.datetime, FlyteDirectory and FlyteFile
   4. I will strip white space because it will be too complicated to support string with white space.
   (datetime.datetime will be affected)
4. I've supported FlyteDirectroy and FlyteFile as input, however, in remote cases, flytecopilot will fail, so I think I can support it locally first, and create an issue to support remote cases.
5. Support direct file paths and template-style references as inputs.
   For example, {{.inputs.infile}} and /var/inputs/infile both are is valid inputs.

How was this patch tested?

1. local execution with 4 examples.
   Reference: https://docs.flyte.org/en/latest/user_guide/customizing_dependencies/raw_containers.html#raw-container
   Note: julia image has error when using arm64, I've tested it by docker run -it into the container and found this error.

2. local execution with 5 prmitive types as input and output.
3. FlyteFile task with input and output
4. FlyteFile task with output only
5. FlyteDirectory task with input and output
6. FlyteDirectory task with output only

Note: I've provided an image on docker hub, you can switch to this branch and test it directly.
image: futureoutlier/rawcontainer:0320

Setup process

git clone https://github.com/flyteorg/flytekit
gh pr checkout 2258
make setup
pip install -e .

python raw_container_local_execution.py

python example:

import logging
from typing import Tuple, List
import datetime
from flytekit import ContainerTask, kwtypes, workflow, task
from flytekit.types.file import FlyteFile
from flytekit.types.directory import FlyteDirectory


logger = logging.getLogger(__file__)

@workflow
def primitive_types(a: int, b: bool, c: float, d: str, e: datetime.datetime, f: datetime.timedelta) \
                            -> Tuple[int, bool, float, str, datetime.datetime, datetime.timedelta]:
    return python_return_same_values(a=a, b=b, c=c, d=d, e=e, f=f)

python_return_same_values = ContainerTask(
    name="python_return_same_values",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs=kwtypes(a=int, b=bool, c=float, d=str, e=datetime.datetime, f=datetime.timedelta),
    outputs=kwtypes(a=int, b=bool, c=float, d=str, e=datetime.datetime, f=datetime.timedelta),
    image="futureoutlier/rawcontainer:0320",
    command=[
        "python",
        "return_same_value.py",
        "{{.inputs.a}}",
        "{{.inputs.b}}",
        "{{.inputs.c}}",
        "{{.inputs.d}}",
        "{{.inputs.e}}",
        "{{.inputs.f}}",
        "/var/outputs",
    ],
)
 

flyte_file_io = ContainerTask(
    name="flyte_file_io",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs=kwtypes(inputs=FlyteFile),
    outputs=kwtypes(out=FlyteFile),
    image="futureoutlier/rawcontainer:0320",
    command=[
        "python",
        "write_flytefile.py",
        "{{.inputs.inputs}}",
        # "/var/inputs/inputs",
        "/var/outputs/out",
    ],
)

flyte_dir_io = ContainerTask(
    name="flyte_dir_io",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    inputs=kwtypes(inputs=FlyteDirectory),
    outputs=kwtypes(out=FlyteDirectory),
    image="futureoutlier/rawcontainer:0320",
    command=[
        "python",
        "write_flytedir.py",
        "{{.inputs.inputs}}",
        # "/var/inputs/inputs",
        "/var/outputs/out",
    ],
)

flyte_dir_out_only = ContainerTask(
    name="flyte_dir_out_only",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    # inputs=kwtypes(inputs=FlyteDirectory),
    outputs=kwtypes(out=FlyteDirectory),
    image="futureoutlier/rawcontainer:0320",
    command=[
        "python",
        "return_flytedir.py",
        "/var/outputs/out",
    ],
)

flyte_file_out_only = ContainerTask(
    name="flyte_file_out_only",
    input_data_dir="/var/inputs",
    output_data_dir="/var/outputs",
    # inputs=kwtypes(inputs=FlyteDirectory),
    outputs=kwtypes(out=FlyteFile),
    image="futureoutlier/rawcontainer:0320",
    command=[
        "python",
        "return_flytefile.py",
        "/var/outputs/out",
    ],
)

@task
def flyte_file_task() -> FlyteFile:
    with open("./a.txt", "w") as file:
        file.write("This is a.txt file.")
    return FlyteFile(path="./a.txt")

@workflow
def flyte_file_io_wf() -> FlyteFile:
    ff = flyte_file_task()
    return flyte_file_io(inputs=ff)

@task
def flyte_dir_task() -> FlyteDirectory:
    from pathlib import Path
    import flytekit
    import os

    working_dir = flytekit.current_context().working_directory
    local_dir = Path(os.path.join(working_dir, "csv_files"))
    local_dir.mkdir(exist_ok=True)
    write_file = local_dir / "a.txt"
    with open(write_file, "w") as file:
        file.write("This is for flyte dir.")

    return FlyteDirectory(path=str(local_dir))

@workflow
def flyte_dir_io_wf() -> FlyteDirectory:
    fd = flyte_dir_task()
    return flyte_dir_io(inputs=fd)

if __name__ == "__main__":
    print(flyte_dir_io_wf())
    print(flyte_file_io_wf())
    print(primitive_types(a=0, b=False, c=3.0, d="hello", e=datetime.datetime.now(), 
                          f=datetime.timedelta(days=1, hours=3, minutes=2, seconds=3, microseconds=5)))
    print(flyte_dir_out_only())
    print(flyte_file_out_only())

Dockerfile

# Use the Alpine Linux version of Python 3.9 as the base image
FROM python:3.9-alpine

# Set the working directory in the container
WORKDIR /root

# You can use the docker copy command after building the image to copy test.py into the /app directory of the container
COPY ./write_flytefile.py /root/write_flytefile.py
COPY ./write_flytedir.py /root/write_flytedir.py
COPY ./return_same_value.py /root/return_same_value.py
COPY ./return_flytefile.py /root/return_flytefile.py
COPY ./return_flytedir.py /root/return_flytedir.py


# Specify the command to run when the container starts. Here, we use /bin/sh to start a shell.
CMD ["/bin/sh"]

write_flytefile.py

import sys
from pathlib import Path

def copy_content_to_output(input_path: Path, output_path: Path):

    content = input_path.read_text()

    output_path.write_text(content)

if __name__ == "__main__":
    if len(sys.argv) > 2:
        input_path = Path(sys.argv[1])
        output_path = Path(sys.argv[2])
        copy_content_to_output(input_path, output_path)
    else:
        print("Usage: script.py <input_path> <output_path>")

write_flytedir.py

import sys
from pathlib import Path
import shutil

def copy_directory(input_path: Path, output_path: Path):
    if not input_path.exists() or not input_path.is_dir():
        print(f"Error: {input_path} does not exist or is not a directory.")
        return

    try:
        shutil.copytree(input_path, output_path)
        print(f"Directory {input_path} was successfully copied to {output_path}")
    except Exception as e:
        print(f"Error copying {input_path} to {output_path}: {e}")

if __name__ == "__main__":
    if len(sys.argv) > 2:
        input_path = Path(sys.argv[1])
        output_path = Path(sys.argv[2])
        copy_directory(input_path, output_path)
    else:
        print("Usage: script.py <input_directory_path> <output_directory_path>")

return_same_value.py

import sys

def write_output(output_dir, output_file, v):
    with open(f"{output_dir}/{output_file}", "w") as f:
            f.write(str(v))

def main(*args, output_dir):
    # Generate output files for each input argument
    for i, arg in enumerate(args, start=1):
        # Using i to generate filenames like 'a', 'b', 'c', ...
        output_file = chr(ord('a') + i - 1)
        write_output(output_dir, output_file, arg)

if __name__ == "__main__":
    *inputs, output_dir = sys.argv[1:]  # Unpack all inputs except for the last one for output_dir

    main(*inputs, output_dir=output_dir)

return_flytefile.py

import sys
from pathlib import Path


def write_output(out_path: Path):
    out_path.write_text("return flyte file content")


if __name__ == "__main__":
    write_output(Path(sys.argv[1]))

return_flytedir.py

import sys
from pathlib import Path

def write_output(directory_path: Path):
    directory_path.mkdir(parents=True, exist_ok=True)
    
    file_path = directory_path / "test_flyte_dir_file"
    
    file_path.write_text("prove flytedir works")

if __name__ == "__main__":
    if len(sys.argv) > 1:
        write_output(Path(sys.argv[1]))
    else:
        print("Usage: script.py <output_directory_path>")

Screenshots

Support documentation example

All examples above

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

#1745

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 85.69%. Comparing base (bf38b8e) to head (a9b0e1e).
Report is 6 commits behind head on master.

❗ Current head a9b0e1e differs from pull request most recent head b068357. Consider uploading reports for the commit b068357 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2258      +/-   ##
==========================================
+ Coverage   83.04%   85.69%   +2.65%     
==========================================
  Files         324       20     -304     
  Lines       24861     1279   -23582     
  Branches     3547        0    -3547     
==========================================
- Hits        20645     1096   -19549     
+ Misses       3591      183    -3408     
+ Partials      625        0     -625     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[kumare3]

Can we check if docker is installed and then use it please do not add docker dependency

[Future-Outlier]

Can we check if docker is installed and then use it please do not add docker dependency

No problem, will do that.

[Future-Outlier]

cc @eapolinario , @pingsutw Could you please help review?
I think I might miss some edge cases.
I will add a unit test after you think is good, thank you!

[Future-Outlier]

cc @wild-endeavor If possible, please take a look too, thank you!

[Future-Outlier]

Question:

1. can I add it in dev-requirements.in? so that we can add unit tests for local execution.
2. are we comfortable with def local_execute's implementation?
   I think it will be better to split it into small functions, however, most of implementation is referenced from here.
   

   flytekit/flytekit/core/base_task.py

   Lines 251 to 328 in 329aea8

   	def local_execute(
   	self, ctx: FlyteContext, **kwargs
   	) -> Union[Tuple[Promise], Promise, VoidPromise, Coroutine, None]:
   	"""
   	This function is used only in the local execution path and is responsible for calling dispatch execute.
   	Use this function when calling a task with native values (or Promises containing Flyte literals derived from
   	Python native values).
   	"""
   	# Unwrap the kwargs values. After this, we essentially have a LiteralMap
   	# The reason why we need to do this is because the inputs during local execute can be of 2 types
   	# - Promises or native constants
   	# Promises as essentially inputs from previous task executions
   	# native constants are just bound to this specific task (default values for a task input)
   	# Also along with promises and constants, there could be dictionary or list of promises or constants
   	try:
   	kwargs = translate_inputs_to_literals(
   	ctx,
   	incoming_values=kwargs,
   	flyte_interface_types=self.interface.inputs,
   	native_types=self.get_input_types(), # type: ignore
   	)
   	except TypeTransformerFailedError as exc:
   	msg = f"Failed to convert inputs of task '{self.name}':\n {exc}"
   	logger.error(msg)
   	raise TypeError(msg) from exc
   	input_literal_map = _literal_models.LiteralMap(literals=kwargs)
   	# if metadata.cache is set, check memoized version
   	local_config = LocalConfig.auto()
   	if self.metadata.cache and local_config.cache_enabled:
   	# TODO: how to get a nice `native_inputs` here?
   	logger.info(
   	f"Checking cache for task named {self.name}, cache version {self.metadata.cache_version} "
   	f"and inputs: {input_literal_map}"
   	)
   	if local_config.cache_overwrite:
   	outputs_literal_map = None
   	logger.info("Cache overwrite, task will be executed now")
   	else:
   	outputs_literal_map = LocalTaskCache.get(self.name, self.metadata.cache_version, input_literal_map)
   	# The cache returns None iff the key does not exist in the cache
   	if outputs_literal_map is None:
   	logger.info("Cache miss, task will be executed now")
   	else:
   	logger.info("Cache hit")
   	if outputs_literal_map is None:
   	outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)
   	# TODO: need `native_inputs`
   	LocalTaskCache.set(self.name, self.metadata.cache_version, input_literal_map, outputs_literal_map)
   	logger.info(
   	f"Cache set for task named {self.name}, cache version {self.metadata.cache_version} "
   	f"and inputs: {input_literal_map}"
   	)
   	else:
   	# This code should mirror the call to `sandbox_execute` in the above cache case.
   	# Code is simpler with duplication and less metaprogramming, but introduces regressions
   	# if one is changed and not the other.
   	outputs_literal_map = self.sandbox_execute(ctx, input_literal_map)
   	if inspect.iscoroutine(outputs_literal_map):
   	return outputs_literal_map
   	outputs_literals = outputs_literal_map.literals
   	# TODO maybe this is the part that should be done for local execution, we pass the outputs to some special
   	# location, otherwise we dont really need to right? The higher level execute could just handle literalMap
   	# After running, we again have to wrap the outputs, if any, back into Promise objects
   	output_names = list(self.interface.outputs.keys()) # type: ignore
   	if len(output_names) != len(outputs_literals):
   	# Length check, clean up exception
   	raise AssertionError(f"Length difference {len(output_names)} {len(outputs_literals)}")
   	# Tasks that don't return anything still return a VoidPromise
   	if len(output_names) == 0:
   	return VoidPromise(self.name)
   	vals = [Promise(var, outputs_literals[var]) for var in output_names]
   	return create_task_output(vals, self.python_interface)

3. Are there any edge cases for container tasks?
   I assume every input should be have prefix "{{.inputs." and suffix "}}"

[eapolinario]

First of all, thanks for you contribution! This will be a massive improvement to the local flytekit use cases.

This is shaping up to be a great PR. I left some comments along the way. But can you make sure to write tests where container tasks are involved in workflows where a combination of regular tasks and container tasks?

[eapolinario]

Skip the tests using pytest.mark.skipif, e.g. . You can combine the expression like @pytest.mark.skipif(sys.platform in ["darwin", "win32"], reason="Skip if running on windows or macos").

[eapolinario]

let's use a regex to handle this. It's a little weird that we're using go templates in python, but since we're not going to do any crazy expressions (it'll be mostly used to parse inputs), something like this will suffice:

    regex = r"^\{\{\s*\.inputs\.(.*?)\s*\}\}$"
    match = re.match(regex, cmd)
    if match:
        return match.group(1)

Notice that since those are go templates, there might be whitespaces before and after the argument inside {{ and }}.

[eapolinario]

can you write a test case to cover this case?

[Future-Outlier]

No problem, will do it

[eapolinario]

not all images have sh installed, right? Why can't we pass the command directly?

[Future-Outlier]

Nice advice, we should pass the command directly.

[eapolinario]

let's make commands an array and use a join (i.e. " ".join(commands)) at the end to build the actual command.

[Future-Outlier]

nice advice, thank you

[Future-Outlier]

update: will use an array as inputs, this could fix cases we need to split strings.

[eapolinario]

we can use remove=True in the invocation of run instead.

[Future-Outlier]

nice advice, thank you

[eapolinario]

this is not neede as the previous call (i.e. the call to wait) only returns after the container stops.

[Future-Outlier]

nice catch, thank you

[Future-Outlier]

First of all, thanks for you contribution! This will be a massive improvement to the local flytekit use cases.

This is shaping up to be a great PR. I left some comments along the way. But can you make sure to write tests where container tasks are involved in workflows where a combination of regular tasks and container tasks?

Really great advice!
I'll do it today, thank you!

[Future-Outlier]

@eapolinario

These are my updates:

1. Command inputs now use an array, and this can make str don't need to be splitter.
   (sh -c is removed now.)
2. Added support for datetime.timedelta.
3. Implemented regex for input parsing.
4. Updated all related tests.
5. Improved string to boolean conversion: output_dict[k] = False if output_val.lower() == "false" else True.

Thank you very much!

[Future-Outlier]

Hi, @eapolinario
Can you help approve this PR?
Companies like AXIS and Tesla might want this feature, thank you!

[Future-Outlier]

This is looking pretty good, but we should increase the test coverage.

Can you add a separate test suite (call it test_local_raw_container.py under tests/flytekit/unit/core) and add a Dockerfile there and use it to build a local image to be used in tests. Add a few more tests to show that the manipulation of input/output data dir works. Also, make sure to add examples where FlyteFile are used. This will help the implementation to catch up with the current state of copilot.

Hi, @eapolinario ,

I've implemented your suggestions:

• Added 3 Python files for testing local execution of FlyteFile, FlyteDirectory, and primitive types within container tasks.
• Created a Dockerfile to include the above test files.
• Introduced 4 tests:
  1. test_flytefile_wf for FlyteFile I/O,
  2. test_flytedir_wf for FlyteDirectory I/O,
  3. test_primitive_types_wf for various primitive types,
  4. test_input_output_dir_manipulation to validate input/output directory handling.

Note: Despite FlyteDirectory not being supported by copilot, I've included it for local execution, ready to extend support based on user feedback.

If it looks good to you, can you help approve it?
This is a huge improvement to companies using ContainerTask, thank you!

[pingsutw]

could we remove _get_key_from_cmd here, and add another function render_command?

def render_command(cmd: str):
    for k, v in self.interface.inputs.items():
        if type(native_inputs[k]) in [FlyteFile, FlyteDirectory]:
            ...
        else:
            cmd = cmd.replace("{{.inputs." + k + "}}", str(native_inputs[k]))

command = [render_command(cmd) for cmd in cmd_and_args]

[Future-Outlier]

amazing advice, thank you.

[eapolinario]

So clean. Thank you!

[Future-Outlier]

So clean. Thank you!

Thank you both too😭😭😭