Data Caching and Hashing

Data Caching and Hashing in flyte-sdk

flyte-sdk leverages data hashing to enable robust task caching and data versioning, optimizing workflow execution by avoiding redundant computations. This mechanism is primarily integrated into the File and Dir classes, which serve as the fundamental interfaces for handling data paths, whether local or remote.

File and Dir: Core Data Abstractions

The File class represents a generic file with a specified format, while the Dir class represents a directory containing files of a specified format. Both provide asynchronous and synchronous interfaces for common I/O operations like opening, downloading, and checking existence. For instance, File offers open() and open_sync() for reading/writing file contents, and Dir provides walk() and walk_sync() to iterate over its contained files.

These classes include path, name, and format attributes to describe the data. Crucially, they also feature hash and hash_method attributes, which are central to flyte-sdk's caching and versioning capabilities.

Data Hashing for Optimized Workflows

Data hashing in flyte-sdk allows the system to determine if the underlying data for a File or Dir object has changed. This is vital for task caching, where a task's output can be reused if its inputs (including data) are identical to a previous run. The hash attribute stores the computed or pre-computed hash value, while hash_method specifies how this hash should be generated.

Providing Pre-computed Hashes

When the hash of a file or directory is already known, you can provide it directly to flyte-sdk using the PrecomputedValue class. This avoids the overhead of recomputing the hash, which is particularly useful for static datasets or when integrating with external systems that provide content hashes.

The PrecomputedValue class implements the HashMethod protocol, but its update method does nothing, and its result method simply returns the value it was initialized with. This makes it suitable for scenarios where the hash is externally managed.

You can supply a pre-computed hash when creating a File or Dir object from an existing remote path, or when uploading a local file:

from flyte.io import File
from flyte.io._hashing_io import PrecomputedValue
from pandas import DataFrame

# Provide a pre-computed hash for an existing remote file
remote_file_with_known_hash = File[DataFrame].from_existing_remote(
    "s3://my-bucket/data.csv", file_cache_key="my_known_hash_value"
)

# Provide a pre-computed hash when uploading a local file
# Note: The actual content of /tmp/data.csv is not hashed by flyte-sdk in this case for cache key calculation.
precomputed_hash_obj = PrecomputedValue("another_known_hash")
remote_file_uploaded = await File[DataFrame].from_local(
    '''/tmp/data.csv''', '''s3://bucket/data.csv''', hash_method=precomputed_hash_obj
)

Similarly, for directories, you can provide a dir_cache_key:

from flyte.io import Dir

remote_dir_with_known_hash = Dir[DataFrame].from_existing_remote(
    "s3://my-bucket/data/", dir_cache_key="dir_hash_abc123"
)

Automatic Hash Computation

flyte-sdk can automatically compute data hashes during file I/O operations. This is achieved by associating a HashMethod implementation with the File object. The HashMethod protocol defines the interface for hash accumulators, requiring update(data: memoryview) to process data chunks and result() -> str to return the final hash.

HashlibAccumulator is a concrete implementation of HashMethod that wraps Python's hashlib library. It can be initialized with a specific hashing algorithm, such as SHA256, and will accumulate the hash as data is read or written.

To enable automatic hashing, you pass an instance of HashlibAccumulator (or any custom HashMethod implementation) to the File object's creation or new_remote method:

import pandas as pd
from flyte.io import File
from flyte.io._hashing_io import HashlibAccumulator
from flytekit.core.annotation import context as env

@env.task
async def write_and_hash_file() -> File[pd.DataFrame]:
    df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
    # Create a new remote file and specify a hash method for automatic computation
    file = File.new_remote(hash_method=HashlibAccumulator.from_hash_name("sha256"))
    async with file.open("wb") as f:
        # As data is written, the HashingWriter will update the hash accumulator
        df.to_csv(f, index=False)
    # The file.hash attribute will now contain the computed SHA256 hash
    print(f"Computed hash: {file.hash}")
    return file

@env.task
async def upload_and_hash_local_file() -> File[pd.DataFrame]:
    # Assume /tmp/data.csv exists
    # Upload a local file, computing its hash during the upload process
    local_path = "/tmp/data.csv"
    # Create a dummy file for demonstration if it doesn't exist
    with open(local_path, "w") as f:
        f.write("col1,col2\n1,a\n2,b\n")

    remote_file = await File[pd.DataFrame].from_local(
        local_path,
        remote_destination="s3://my-bucket/uploaded_data.csv",
        hash_method=HashlibAccumulator.from_hash_name("sha256")
    )
    print(f"Uploaded file hash: {remote_file.hash}")
    return remote_file

Under the Hood: Hashing I/O Wrappers

The actual byte-level hashing during read and write operations is handled by internal I/O wrapper classes within the io._hashing_io package. These classes intercept data as it flows through file handles and pass it to the configured HashMethod accumulator.

• HashingWriter and AsyncHashingWriter: These classes wrap synchronous and asynchronous file handles, respectively. When data is written through them, they update the associated HashMethod instance before passing the data to the underlying file handle.
• HashingReader and AsyncHashingReader: Similarly, these classes wrap synchronous and asynchronous file handles for reading. As data is read, they update the HashMethod instance with the read bytes.

Users typically do not interact directly with these wrapper classes. Instead, they are transparently employed by the File class when a hash_method is specified for operations like File.open() (for writing remote files) or File.from_local() (for uploading local files).

Gotchas and Best Practices

• Synchronous Writes and Hashing: The File.open_sync() method does not currently integrate with the hashing logic. If you perform synchronous writes and require automatic hash computation, you must ensure the hash is computed by other means or use the asynchronous File.open() method with a hash_method.
• Binary Mode for Remote Files: When opening remote files for writing using File.open(), the mode string must include `