fix: Mark flaky tests to rerun on failure

[bhimrazy]

What does this PR do ?

Mark specific tests as flaky to allow them to rerun automatically upon failure, improving test reliability.

Partially covers: #437

Functions falling under this:

• tests/streaming/test_dataset.py::test_subsample_streaming_dataset_with_token_loader
• tests/streaming/test_dataset.py::test_dataset_for_text_tokens_distributed_num_workers_end_to_end

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 80%. Comparing base (1a0f48c) to head (47333f7).
Report is 1 commits behind head on main.

Additional details and impacted files

@@         Coverage Diff         @@
##           main   #610   +/-   ##
===================================
  Coverage    80%    80%           
===================================
  Files        43     43           
  Lines      6463   6463           
===================================
  Hits       5156   5156           
  Misses     1307   1307           

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.