Streaming Literals

[cosmicBboy]

Streaming Literals: Sharing partial state between tasks via a publisher/subscriber model.

Author: @kumare3

Intro

Today in Flyte we have Literals. A Literal has an implicit contract - it is materialized when the producer task produces it and the consumer task cannot run, unless it is fully materialized. But there are cases in which, we may want to to partially materialize or not at all materialize a Literal, but start consuming it in the partial state. This can be generalized to a stream and a publisher/subscriber model.

Why this might be useful?

• High-speed training with better gpu utilization. Or also avoid materializing transform outputs and using them in the training.
• We could build a full end to end Streaming pipeline composed of disparate streaming compute engines with this model
• This enables container re-use, with loss of recovery per data point, but separation of compute units.

Code sample

from flytekit import task, workflow, Stream

@task
def foo1(i: int, j: Stream[int]) -> Stream[int]:
   for v in j:
      yield v * i

@task
def foo2(k: Stream[int]) -> Stream[float]:
    for i in k:
        yield float(i)
    
@workflow
def streamer(i: int, j: Stream[int]) -> Stream[float]:
    yield foo2(k=foo1(i=i, j=j))

Example

Let us use an example to motivate the usecase. Consider a very large (100 billion ints) array of elements, for which the subsequent task wants to compute the sum of all the numbers.

flowchart LR
    A[Task A] -- number --> B[Sum of all numbers]

Loading

This can be done in 3 ways:

Approach 1: Current Predominant Flyte Pattern

• Task A Generates all the numbers once, materialize them in S3
• Task B downloads the numbers from S3 and computes the sum

Approach 2: Supported by Flyte

• Task A Generates all the numbers once, materialize them in S3
• Task B streams the numbers from S3 and computes the sum

Approach 3: Not possible with current Flyte compiler and Propeller model

• Task A Generates all the numbers once and materializes them to a disk / queue and opens an endpoint for other consumers to read from
• Task B streams data from the queue / directly from Task A and computes a sum. Finally outputs the sum

Analysis

• Approach 1 is the least efficient, but has the best recoverability.
• Approach 2 is more efficient, and can have similar recoverability and checkpointing can be used to improve the recovery process
• Approach 3 is the most efficient, but failure is harder to handle.

RFC proposal

Approach 3 could make it possible to build partial streaming applications in Flyte. To achieve this one option could be to build a new type called StreamLiteral

Every task that produces a streamLiteral and subsequent task that consumes the streamLiteral have a soft dependency, where they can be launched together. The compiler adds a streamID to every streamLiteral. The streamLiteral itself can hold a Literal as the datum carried by the Stream.

At runtime, the stream ID could be materialized using multiple options,

1. A service on the producer task
2. A buffer / queue like NATS, Kafka between the producer and the consumer task
3. a Shared filesystem with the producer task as the writer and the consumer task as the reader.

A few constraints on a such a system could be,

1. Restrict producers to be one task
2. Producer task can be arbitrarily complex - can be a streaming Flink Job
3. Consumers maintain their own state.

In a static graph, all consumers are known so the stream will not be closed until all consumers exit.

[fg91]

Very interesting feature imho, was asked by users whether something like this is possible in Flyte before. Let's turn this into an RFC 👍

[AdemFr]

I was the one asking, and yes it would be amazing. Some use cases that are interesting:

• transcoding of many videos
• patching images into multiple ones and immediately process the patches per image (basically all one to many map type operations)

[davidmirror-ops]

2023-11-09 Contributor's meetup notes: this entry meets the requirements for an RFC; @cosmicBboy feel free to write a new RFC.

[cosmicBboy]

Would that be a markdown file here? https://github.com/flyteorg/flyte/tree/master/rfc/core%20language

[davidmirror-ops]

@cosmicBboy yes :)
https://github.com/flyteorg/flyte/tree/master/rfc#the-process

[davidmirror-ops]

@wild-endeavor is this connected to the Raw Literals Offload work?