def test_serialization_settings_transport():
default_img = Image(name="default", fqn="test", tag="tag")
serialization_settings = SerializationSettings(
project="project",
domain="domain",
version="version",
env={"hello": "blah"},
image_config=ImageConfig(
default_image=default_img,
images=[default_img],
),
flytekit_virtualenv_root="/opt/venv/blah",
python_interpreter="/opt/venv/bin/python3",
fast_serialization_settings=FastSerializationSettings(
enabled=True,
destination_dir="/opt/blah/blah/blah",
distribution_location="s3://my-special-bucket/blah/bha/asdasdasd/cbvsdsdf/asdddasdasdasdasdasdasd.tar.gz",
),
)
tp = serialization_settings.serialized_context
with_serialized = serialization_settings.with_serialized_context()
assert serialization_settings.env == {"hello": "blah"}
assert with_serialized.env
assert with_serialized.env[SERIALIZED_CONTEXT_ENV_VAR] == tp
ss = SerializationSettings.from_transport(tp)
assert ss is not None
assert ss == serialization_settings
assert len(tp) == 376
def get_serializable_task(
entity_mapping: OrderedDict,
settings: SerializationSettings,
entity: FlyteLocalEntity,
) -> task_models.TaskSpec:
task_id = _identifier_model.Identifier(
_identifier_model.ResourceType.TASK,
settings.project,
settings.domain,
entity.name,
settings.version,
)
if isinstance(
entity, PythonFunctionTask
) and entity.execution_mode == PythonFunctionTask.ExecutionBehavior.DYNAMIC:
# In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state
# from the serialization context. This is passed through an environment variable, that is read from
# during dynamic serialization
settings = settings.with_serialized_context()
container = entity.get_container(settings)
# This pod will be incorrect when doing fast serialize
pod = entity.get_k8s_pod(settings)
if settings.should_fast_serialize():
# This handles container tasks.
if container and isinstance(entity,
(PythonAutoContainerTask, MapPythonTask)):
# For fast registration, we'll need to muck with the command, but on
# ly for certain kinds of tasks. Specifically,
# tasks that rely on user code defined in the container. This should be encapsulated by the auto container
# parent class
container._args = prefix_with_fast_execute(settings,
container.args)
# If the pod spec is not None, we have to get it again, because the one we retrieved above will be incorrect.
# The reason we have to call get_k8s_pod again, instead of just modifying the command in this file, is because
# the pod spec is a K8s library object, and we shouldn't be messing around with it in this file.
elif pod:
if isinstance(entity, MapPythonTask):
entity.set_command_prefix(
get_command_prefix_for_fast_execute(settings))
pod = entity.get_k8s_pod(settings)
else:
entity.set_command_fn(
_fast_serialize_command_fn(settings, entity))
pod = entity.get_k8s_pod(settings)
entity.reset_command_fn()
tt = task_models.TaskTemplate(
id=task_id,
type=entity.task_type,
metadata=entity.metadata.to_taskmetadata_model(),
interface=entity.interface,
custom=entity.get_custom(settings),
container=container,
task_type_version=entity.task_type_version,
security_context=entity.security_context,
config=entity.get_config(settings),
k8s_pod=pod,
sql=entity.get_sql(settings),
)
if settings.should_fast_serialize() and isinstance(
entity, PythonAutoContainerTask):
entity.reset_command_fn()
return task_models.TaskSpec(template=tt)