Enable Ray Fast Register

[fiedlerNr9]

Why are the changes needed?

Today, when launching a Ray job in Flyte only the Head Node is retrieving fast register data - no applying for Ray Worker Nodes. This results in the need for the user of copy the source code into the containers working dir & rebuilding the container image for every change on the ray user code.

This PR utilizes rays functionality of adding files to the ray job on runtime passing the runtime_env {"working_dir":"some_path"}.

⚠️ There is an ray gcs storage limit of 500MiB. Thats why ray supports adding an excludes variable along with working_dir which functions as a .ignore file. This PR introduces a excludes_working_dir to the RayJobConfig, which allows control over what files to upload to the ray cluster.

What changes were proposed in this pull request?

• Adding working_dir & excludes init parameter when running remotely
• Docs on what this parameter is doing
• TLDR: after head node downloads fast register data -> ray.init("working_dir"="/root", excludes=["data"]) zips all fast register data, takes excludes into account and uploads to ray gcs -> ray worker nodes downloads zips from ray gcs and unzips

How was this patch tested?

• Running the following workflow with Imagespec works:

from flytekit import ImageSpec, Resources, task, workflow, PodTemplate
from flytekitplugins.ray import HeadNodeConfig, RayJobConfig, WorkerNodeConfig
from ray.util.multiprocessing import Pool

import time
import math
import random

container_image = ImageSpec(
    name="rayworkerlove",
    apt_packages=["git", "wget"],
    packages=[
        "git+https://github.com/flyteorg/flytekit.git@d5549cee2a9bef8338bd9e94aa23397631386c10#subdirectory=plugins/flytekit-ray"
    ],
    registry="ghcr.io/fiedlerNr9",
)


def sample(num_samples):
    num_inside = 0
    for _ in range(num_samples):
        x, y = random.uniform(-1, 1), random.uniform(-1, 1)
        if math.hypot(x, y) <= 1:
            num_inside += 1
    return num_inside


ray_config = RayJobConfig(
    head_node_config=HeadNodeConfig(
        ray_start_params={"log-color": "True", "num-cpus": "1"}
    ),
    worker_node_config=[
        WorkerNodeConfig(
            group_name="ray-group", replicas=3, min_replicas=1, max_replicas=4
        )
    ],
    shutdown_after_job_finishes=True,
    ttl_seconds_after_finished=120,
    enable_autoscaling=False,
    excludes_working_dir=["micromamba"],
)


@task(
    task_config=ray_config,
    requests=Resources(mem="2Gi", cpu="1"),
    container_image=container_image,
)
def approximate_pi_distributed(num_samples: int):

    pool = Pool()

    start = time.time()
    num_inside = 0
    sample_batch_size = 100000
    for result in pool.map(
        sample, [sample_batch_size for _ in range(num_samples // sample_batch_size)]
    ):
        num_inside += result

    print("pi ~= {}".format((4 * num_inside) / num_samples))
    print("Finished in: {:.2f}s".format(time.time() - start))


@workflow
def wf(n: int = 60000000):
    approximate_pi_distributed(num_samples=n)

Setup process

Screenshots

Check all the applicable boxes

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

[thomasjpfan]

This depends on the system having touch + find, which is usually the case.

To be safer, can this use os.utime + os.walk (or Path.rglob) to update the mtimes?

[fiedlerNr9]

Good catch! I can try this out & adjust. Also this should only stay short term. @pingsutw mentioned there is work in progress that allows us to not set mtime=0 for the fast register tar.gz

[kumare3]

i do not think we will do this right now @fiedlerNr9, because the mtime=0 makes it possible to get consistent hashes for the similar tar files, otherwise we get multiple uploads to admin.

I think you should do what @thomasjpfan recommends use os.walk. the os.system will break and will cause random bugs.

[fiedlerNr9]

No need to manually adjust the timestamp for this now. I just catch the fast register data and dont zip it at all

[thomasjpfan]

Can this comment say explicitly how the data is being moved around? My understanding is:

1. Flyte's fast register and sends it to the head node
2. Ray then zips the working directory and sends it to the ray gcs
3. Ray worker pulls the files from ray's gcs.

Is this correct?

[fiedlerNr9]

Its exactly this only with the addition that ray respects the given excludes variable (list[str]), which functions as a ignore file

[pingsutw]

qq: could we exclude the file during the fast registration? (.gitignore)
Do we have a case where we want to upload a file to the head node but don't want it on the worker nodes?"

[fiedlerNr9]

Do we have a case where we want to upload a file to the head node but don't want it on the worker nodes?

No i dont think so.

What file are you trying to exclude? The fast register .tar.gz? Actually I am thinking, its only the .tar.gz file coming with the mtime=0 right? If we exclude that we dont need to modify any timestamp on files?

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 78.56%. Comparing base (eae31c0) to head (061698f).
Report is 9 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #2606       +/-   ##
===========================================
+ Coverage   47.17%   78.56%   +31.38%     
===========================================
  Files         230      185       -45     
  Lines       21322    18881     -2441     
  Branches     3711     3714        +3     
===========================================
+ Hits        10059    14834     +4775     
+ Misses      11154     3376     -7778     
- Partials      109      671      +562     

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

[fiedlerNr9]

Okay no need to manually update timestamps here. I am just always put the in the exclude variable, since there is no need to upload the .tar.gz to the ray gcs.

"excludes": ["script_mode.tar.gz", "fast*.tar.gz"]!

Works and tested with pyflyte run and register. Would appreciate another review

[eapolinario]

Jus to reiterate, this looks great. Thanks, @fiedlerNr9 !