Fix flaky E2E test in Trainer Notebook examples

[andreyvelich]

Some of our E2E tests are flaky and fails time to time: https://github.com/kubeflow/trainer/actions/workflows/test-e2e.yaml
We should explore how to improve Notebook example to resolve it.
Example: https://github.com/kubeflow/trainer/actions/runs/23232456393

/area testing
/good-first-issue

[google-oss-prow]

@andreyvelich:
This request has been marked as suitable for new contributors.

Please ensure the request meets the requirements listed here.

If this request no longer meets these requirements, the label can be removed
by commenting with the /remove-good-first-issue command.

Details

In response to this:

Some of our E2E tests are flaky and fails time to time: https://github.com/kubeflow/trainer/actions/workflows/test-e2e.yaml
We should explore how to improve Notebook example to resolve it.
Example: https://github.com/kubeflow/trainer/actions/runs/23232456393

/area testing
/good-first-issue

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

[Amir380-A]

Hi @andreyvelich
I would like to work on this.

[neeraj542]

Hey @andreyvelich I would be happy to work on this issue, could you pls assign this to me

[xikronz]

/assign

[neeraj542]

/assign

[xikronz]

Hi @andreyvelich! following up on about what we briefly talked about on call, I don't think the problem is a timeout issue at all. I looked into a the artifacts of a couple failed tests e2e 1.32.3, e2e 1.34.0, e2e 1.35.0, another e2e 1.32.3 etc.. and they're all triggered by the same cell [9] in ./examples/local/local-container-mnist.ipynb around 1 minute into the execution process which is the MNIST dataset download step, note that the progress bar is always at 100% in these ipynb cells so it seems reasonable to start debugging from this checkpoint.

If we follow the error outputs in the artifacts files from cell number 8 a bit more closely, you'll notice they're all variant's of this error

which comes from the torchvision package. it's possible to trace the exact runtime error back to this function in their source code [line 137 in function 87]:

def download_url(url, root, filename=None, md5=None):
    # ... code for downloads file ...
    if md5 is not None and not check_md5(fpath, md5):
        raise RuntimeError("File not found or corrupted.")

So it seems like the AWS download endpoint being used for the dataset intermittently returns content that fails the expected checksum from within the GitHub Actions environment (even though it returns HTTP 200 and the download progress reaches 100%). The following dist.barrier() failures (e.g. “Connection closed by peer”) for the second container in the same output cell

which is consistent with this idea that one rank is crashing early and the other rank hitting the collective after the peer has exited, i.e. a cascade from the original dataset integrity error.

Also, if we dig into the SDK backend.py file and find the client.wait_for_job_status() function on line 635, you'll notice that if it was a timeout error we'd raise a raise TimeoutError(f"Timeout waiting for TrainJob {name} to reach status: {status}") but every single failed E2E test is rooted in a RuntimeError: TrainJob <JOB_ID> is Failed error so I don't think increasing timeout to 300 will fix anything

If you agree that this is the root cause, I think there's two ways to proceed:

1. switch our download strategy (from datasets.FashionMNIST to datasets.MNIST)
2. add retry / richer logging / prefetching so the artifacts clearly show the underlying URL+MD5 mismatch when it happens again

Let me know your thoughts!!

[xikronz]

TLDR: every E2E test that fails ends up being traced back to local-container-mnist.ipynb E2E and it seems like they're caused by a MD5 checksum failure when downloading FashionMNIST from the AWS endpoint inside GitHub Actions. Fix options: switch to a more reliable dataset source, or add retry logic + better logging around the download step!

[xikronz]

Also while we're discussing these notebooks, there's a race condition occurring in the download step that could be unsafe when nodes>1 (which is the case for these container_mnist notebooks). Notice how in both pictures I attached there's a local_rank=0 flag occurring at the very end despite one coming from [rank0] and the other rank1. This means they both satisfy this loop guard condition: if local_rank == 0: dataset = datasets.FashionMNIST(...)

This wouldn't be an issue if Kubeflow loads a seperate backend for each container or if the I/O was super fast but if we dig a little deeper and look into the backend.py file of the SDK, we can see that the train() method creates a single host working directory [line 258/259]:

# Create per-job working directory on host (for outputs, checkpoints, etc.)
workdir = container_utils.create_workdir(trainjob_name)

On line [314/394], the for loop that iterates over the ranks uses the same workdir constant to /workspace with read-write access:

for rank in range(num_nodes):
    # ... master_addr, entry_cmd, labels, etc. ...
      volumes = {    
           workdir: {        
               "bind": constants.WORKSPACE_PATH,  
               "mode": "rw",    }}

This means every container gets the same host directory bound to /workspace so ./data in the training script (under the container’s working dir /workspace) is the same shared directory for all nodes. The concurrent writes from both ranks corrupt the dataset files --> MD5 verification fails: RuntimeError: File not found or corrupted and we crash... The only thing is that this would be more of a deterministic bug instead of a flaky one so I'm less sure about it. Anyways, very easy fix here just do dist.get_rank() == 0 (global rank).