Offload literals

flytekit/bin/entrypoint.py

@@ -12,7 +12,7 @@
 import traceback
 import warnings
 from sys import exit
-from typing import Callable, List, Optional
+from typing import Callable, Dict, List, Optional
 
 import click
 from flyteidl.core import literals_pb2 as _literals_pb2
@@ -137,7 +137,41 @@ def _dispatch_execute(
             logger.warning("Task produces no outputs")
             output_file_dict = {_constants.OUTPUT_FILE_NAME: _literal_models.LiteralMap(literals={})}
         elif isinstance(outputs, _literal_models.LiteralMap):
-            output_file_dict = {_constants.OUTPUT_FILE_NAME: outputs}
+            offloaded_literals: Dict[str, _literal_models.Literal] = {}
+            literal_map_copy = {}
+
+            min_offloaded_size = int(os.environ.get("FK_L_MIN_SIZE_MB", "10")) * 1024 * 1024
+            max_offloaded_size = int(os.environ.get("FK_L_MAX_SIZE_MB", "1000")) * 1024 * 1024
+
+            # Go over each output and create a separate offloaded in case its size is too large
+            for k, v in outputs.literals.items():
+                lit = v.to_flyte_idl()
+                if lit.ByteSize() >= max_offloaded_size:
+                    raise ValueError(
+                        f"Literal {k} is too large to be offloaded. Max literal size is {max_offloaded_size} whereas the literal size is {lit.ByteSize()} bytes"
+                    )
+
+                if lit.ByteSize() >= min_offloaded_size:
+                    logger.debug(f"Literal {k} is too large to be inlined, offloading to metadata bucket")
+
+                    # TODO: hash calculation
+
+                    offloaded_filename = f"{k}_offloaded_metadata.pb"
+
+                    offloaded_literal = _literal_models.Literal(
+                        offloaded_metadata=_literal_models.LiteralOffloadedMetadata(
+                            uri=f"{output_prefix}/{offloaded_filename}",
+                            size_bytes=lit.ByteSize(),
+                            # TODO: do I have to set the inferred literal type?
+                        )
+                    )
+                    literal_map_copy[k] = offloaded_literal
+                    offloaded_literals[offloaded_filename] = v
+                else:
+                    literal_map_copy[k] = v
+            outputs = _literal_models.LiteralMap(literals=literal_map_copy)
+
+            output_file_dict = {_constants.OUTPUT_FILE_NAME: outputs, **offloaded_literals}
         elif isinstance(outputs, _dynamic_job.DynamicJobSpec):
             output_file_dict = {_constants.FUTURES_FILE_NAME: outputs}
         else:

flytekit/interaction/string_literals.py

@@ -61,6 +61,9 @@ def literal_string_repr(lit: Literal) -> typing.Any:
         return [literal_string_repr(i) for i in lit.collection.literals]
     if lit.map:
         return {k: literal_string_repr(v) for k, v in lit.map.literals.items()}
+    if lit.offloaded_metadata:
+        # TODO: load literal from offloaded literal?
+        return f"Offloaded literal metadata: {lit.offloaded_metadata}"
     raise ValueError(f"Unknown literal type {lit}")
 
 

flytekit/models/literals.py

@@ -991,7 +991,7 @@ def to_flyte_idl(self):
             map=self.map.to_flyte_idl() if self.map is not None else None,
             hash=self.hash,
             metadata=self.metadata,
-            offloaded_metadata=self.offloaded_metadata.to_flyte_idl() if self.offloaded_metadata else None,
+            offloaded_metadata=self.offloaded_metadata.to_flyte_idl() if self.offloaded_metadata is not None else None,
         )
 
     @classmethod

tests/flytekit/unit/bin/test_python_entrypoint.py

@@ -1,3 +1,5 @@
+from dataclasses import dataclass
+from datetime import datetime
 import os
 import re
 import textwrap
@@ -8,21 +10,26 @@
 import mock
 import pytest
 from flyteidl.core.errors_pb2 import ErrorDocument
+from flyteidl.core import literals_pb2
+from flyteidl.core.literals_pb2 import Literal, LiteralCollection, Scalar, Primitive
 
 from flytekit.bin.entrypoint import _dispatch_execute, normalize_inputs, setup_execution, get_traceback_str
 from flytekit.configuration import Image, ImageConfig, SerializationSettings
+from flytekit.core import mock_stats
+from flytekit.models.core import identifier as id_models
 from flytekit.core import context_manager
 from flytekit.core.base_task import IgnoreOutputs
 from flytekit.core.dynamic_workflow_task import dynamic
 from flytekit.core.promise import VoidPromise
 from flytekit.core.task import task
 from flytekit.core.type_engine import TypeEngine
 from flytekit.exceptions import user as user_exceptions
-from flytekit.exceptions.scopes import system_entry_point
+from flytekit.exceptions.scopes import system_entry_point, user_entry_point
 from flytekit.exceptions.user import FlyteRecoverableException, FlyteUserRuntimeException
 from flytekit.models import literals as _literal_models
-from flytekit.models.core import errors as error_models
+from flytekit.models.core import errors as error_models, execution
 from flytekit.models.core import execution as execution_models
+from flytekit.core.utils import write_proto_to_file
 
 
 @mock.patch("flytekit.core.utils.load_proto_from_file")
@@ -453,3 +460,253 @@ def test_get_traceback_str():
     expected_error_re = re.compile(expected_error_pattern)
     print(traceback_str)  # helpful for debugging
     assert expected_error_re.match(traceback_str) is not None
+
+
+def test_dispatch_execute_offloaded_literals(tmp_path_factory):
+    @task
+    def t1(a: typing.List[int]) -> typing.List[str]:
+        return [f"string is: {x}" for x in a]
+
+    inputs_path = tmp_path_factory.mktemp("inputs")
+    outputs_path = tmp_path_factory.mktemp("outputs")
+
+    ctx = context_manager.FlyteContext.current_context()
+    with context_manager.FlyteContextManager.with_context(
+            ctx.with_execution_state(
+                ctx.execution_state.with_params(
+                    engine_dir=tmp_path_factory.mktemp("engine_dir"),
+                    mode=context_manager.ExecutionState.Mode.TASK_EXECUTION,
+                    user_space_params=context_manager.ExecutionParameters(
+                        execution_date=datetime.now(),
+                        tmp_dir="/tmp",
+                        stats=mock_stats.MockStats(),
+                        logging=None,
+                        raw_output_prefix="",
+                        output_metadata_prefix=str(outputs_path.absolute()),
+                        execution_id=id_models.WorkflowExecutionIdentifier("p", "d", "n"),
+                    ),
+                ),
+            ),
+    ) as ctx:
+        xs: typing.List[int] = [1, 2, 3]
+        input_literal_map = _literal_models.LiteralMap(
+            {
+                "a": _literal_models.Literal(
+                    collection=_literal_models.LiteralCollection(
+                        literals=[
+                            _literal_models.Literal(
+                                scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x)),
+                            ) for x in xs
+                        ]
+                    )
+                )
+            }
+        )
+
+        write_proto_to_file(input_literal_map.to_flyte_idl(), str(inputs_path/"inputs.pb"))
+
+        with mock.patch.dict(os.environ, {"FK_L_MIN_SIZE_MB": "0"}):
+            _dispatch_execute(ctx, lambda: t1, str(inputs_path/"inputs.pb"), str(outputs_path.absolute()))
+
+            assert "error.pb" not in os.listdir(outputs_path)
+
+            for ff in os.listdir(outputs_path):
+                with open(outputs_path/ff, "rb") as f:
+                    if ff == "outputs.pb":
+                        lit = literals_pb2.LiteralMap()
+                        lit.ParseFromString(f.read())
+                        assert len(lit.literals) == 1
+                        assert "o0" in lit.literals
+                        assert lit.literals["o0"].offloaded_metadata is not None
+                        assert lit.literals["o0"].offloaded_metadata.size_bytes == 62
+                        assert lit.literals["o0"].offloaded_metadata.uri.endswith("/o0_offloaded_metadata.pb")
+                    elif ff == "o0_offloaded_metadata.pb":
+                        lit = literals_pb2.Literal()
+                        lit.ParseFromString(f.read())
+                        assert lit == Literal(
+                            collection=LiteralCollection(
+                                literals=[
+                                    Literal(
+                                        scalar=Scalar(primitive=Primitive(string_value="string is: 1")),
+                                    ),
+                                    Literal(
+                                        scalar=Scalar(primitive=Primitive(string_value="string is: 2")),
+                                    ),
+                                    Literal(
+                                        scalar=Scalar(primitive=Primitive(string_value="string is: 3")),
+                                    ),
+                                ]
+                            )
+                        )
+                    else:
+                        assert False, f"Unexpected file {ff}"
+
+
+def test_dispatch_execute_offloaded_literals_two_outputs_offloaded(tmp_path_factory):
+    @task
+    def t1(xs: typing.List[int]) -> typing.Tuple[int, typing.List[str]]:
+        return sum(xs), [f"string is: {x}" for x in xs]
+
+    inputs_path = tmp_path_factory.mktemp("inputs")
+    outputs_path = tmp_path_factory.mktemp("outputs")
+
+    ctx = context_manager.FlyteContext.current_context()
+    with context_manager.FlyteContextManager.with_context(
+            ctx.with_execution_state(
+                ctx.execution_state.with_params(
+                    engine_dir=tmp_path_factory.mktemp("engine_dir"),
+                    mode=context_manager.ExecutionState.Mode.TASK_EXECUTION,
+                    user_space_params=context_manager.ExecutionParameters(
+                        execution_date=datetime.now(),
+                        tmp_dir="/tmp",
+                        stats=mock_stats.MockStats(),
+                        logging=None,
+                        raw_output_prefix="",
+                        output_metadata_prefix=str(outputs_path.absolute()),
+                        execution_id=id_models.WorkflowExecutionIdentifier("p", "d", "n"),
+                    ),
+                ),
+            ),
+    ) as ctx:
+        xs: typing.List[int] = [1, 2, 3, 4]
+        input_literal_map = _literal_models.LiteralMap(
+            {
+                "xs": _literal_models.Literal(
+                    collection=_literal_models.LiteralCollection(
+                        literals=[
+                            _literal_models.Literal(
+                                scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=x)),
+                            ) for x in xs
+                        ]
+                    )
+                )
+            }
+        )
+
+        write_proto_to_file(input_literal_map.to_flyte_idl(), str(inputs_path/"inputs.pb"))
+
+        with mock.patch.dict(os.environ, {"FK_L_MIN_SIZE_MB": "0"}):
+            _dispatch_execute(ctx, lambda: t1, str(inputs_path/"inputs.pb"), str(outputs_path.absolute()))
+
+            assert "error.pb" not in os.listdir(outputs_path)
+
+            for ff in os.listdir(outputs_path):
+                with open(outputs_path/ff, "rb") as f:
+                    if ff == "outputs.pb":
+                        lit = literals_pb2.LiteralMap()
+                        lit.ParseFromString(f.read())
+                        assert len(lit.literals) == 2
+                        assert "o0" in lit.literals
+                        assert lit.literals["o0"].offloaded_metadata is not None
+                        assert lit.literals["o0"].offloaded_metadata.size_bytes == 6
+                        assert lit.literals["o0"].offloaded_metadata.uri.endswith("/o0_offloaded_metadata.pb")
+                        assert "o1" in lit.literals
+                        assert lit.literals["o1"].offloaded_metadata is not None
+                        assert lit.literals["o1"].offloaded_metadata.size_bytes == 82
+                        assert lit.literals["o1"].offloaded_metadata.uri.endswith("/o1_offloaded_metadata.pb")
+                    elif ff == "o0_offloaded_metadata.pb":
+                        lit = literals_pb2.Literal()
+                        lit.ParseFromString(f.read())
+                        assert lit == Literal(
+                            scalar=Scalar(primitive=Primitive(integer=10)),
+                        )
+                    elif ff == "o1_offloaded_metadata.pb":
+                        lit = literals_pb2.Literal()
+                        lit.ParseFromString(f.read())
+                        assert lit == Literal(
+                            collection=LiteralCollection(
+                                literals=[
+                                    Literal(
+                                        scalar=Scalar(primitive=Primitive(string_value="string is: 1")),
+                                    ),
+                                    Literal(
+                                        scalar=Scalar(primitive=Primitive(string_value="string is: 2")),
+                                    ),
+                                    Literal(
+                                        scalar=Scalar(primitive=Primitive(string_value="string is: 3")),
+                                    ),
+                                    Literal(
+                                        scalar=Scalar(primitive=Primitive(string_value="string is: 4")),
+                                    ),
+                                ]
+                            )
+                        )
+                    else:
+                        assert False, f"Unexpected file {ff}"
+
+
+def test_dispatch_execute_offloaded_literals_two_outputs_only_second_one_offloaded(tmp_path_factory):
+    @dataclass
+    class DC:
+        a: typing.List[int]
+        b: typing.List[str]
+
+    @task
+    def t1(n: int) -> typing.Tuple[int, DC]:
+        return n, DC(a=list(range(n)), b=[f"string is: {x}" for x in range(n)])
+
+    inputs_path = tmp_path_factory.mktemp("inputs")
+    outputs_path = tmp_path_factory.mktemp("outputs")
+
+    ctx = context_manager.FlyteContext.current_context()
+    with context_manager.FlyteContextManager.with_context(
+            ctx.with_execution_state(
+                ctx.execution_state.with_params(
+                    engine_dir=tmp_path_factory.mktemp("engine_dir"),
+                    mode=context_manager.ExecutionState.Mode.TASK_EXECUTION,
+                    user_space_params=context_manager.ExecutionParameters(
+                        execution_date=datetime.now(),
+                        tmp_dir="/tmp",
+                        stats=mock_stats.MockStats(),
+                        logging=None,
+                        raw_output_prefix="",
+                        output_metadata_prefix=str(outputs_path.absolute()),
+                        execution_id=id_models.WorkflowExecutionIdentifier("p", "d", "n"),
+                    ),
+                ),
+            ),
+    ) as ctx:
+        input_literal_map = _literal_models.LiteralMap(
+            {
+                "n": _literal_models.Literal(
+                    scalar=_literal_models.Scalar(primitive=_literal_models.Primitive(integer=56_000)),
+                )
+            }
+        )
+
+        write_proto_to_file(input_literal_map.to_flyte_idl(), str(inputs_path/"inputs.pb"))
+
+        # Notice how the threshold is set to 1MB
+        with mock.patch.dict(os.environ, {"FK_L_MIN_SIZE_MB": "1"}):
+            _dispatch_execute(ctx, lambda: t1, str(inputs_path/"inputs.pb"), str(outputs_path.absolute()))
+
+            assert "error.pb" not in os.listdir(outputs_path)
+
+            # o0 is not offloaded
+            assert "o0_offloaded_metadata.pb" not in os.listdir(outputs_path)
+
+            for ff in os.listdir(outputs_path):
+                with open(outputs_path/ff, "rb") as f:
+                    if ff == "outputs.pb":
+                        lit = literals_pb2.LiteralMap()
+                        lit.ParseFromString(f.read())
+                        assert len(lit.literals) == 2
+
+                        # o0 is not offloaded
+                        assert "o0" in lit.literals
+                        assert lit.literals["o0"].HasField("offloaded_metadata") is False
+
+                        # o1 is offloaded
+                        assert "o1" in lit.literals
+                        assert lit.literals["o1"].HasField("offloaded_metadata") is True
+                        assert lit.literals["o1"].offloaded_metadata.size_bytes == 1108538
+                        assert lit.literals["o1"].offloaded_metadata.uri.endswith("/o1_offloaded_metadata.pb")
+                    elif ff == "o1_offloaded_metadata.pb":
+                        lit = literals_pb2.Literal()
+                        lit.ParseFromString(f.read())
+                        # Load the dataclass from the proto
+                        transformer = TypeEngine.get_transformer(DC)
+                        dc = transformer.to_python_value(ctx, _literal_models.Literal.from_flyte_idl(lit), DC)
+                        assert dc.a == list(range(56_000))
+                    else:
+                        assert False, f"Unexpected file {ff}"