chore: import upstream snapshot with attribution

release/tune_tests/frequent_pausing/script.py

@@ -0,0 +1,58 @@
+"""Run Tune with frequent pausing.
+See context https://github.com/ray-project/ray/issues/34197.
+
+m5.large node has memory of 7.2 GB. With `RAY_memory_usage_threshold=0.5`,
+if the node's memory exceeds 3.6 GB, any new tasks would be killed.
+Note this node memory is also shared by processes like ray dashboard etc.
+Without ray object store reference leakage from application code, all these
+background processes take less than 2 GB of memory all together.
+With reference leakage, we reach 3.6 GB threshold within 5 minutes
+at the time when this test was written.
+
+success criteria: run through 10min without crash.
+
+cost: A few dollars.
+"""
+
+import numpy as np
+import os
+import pickle
+import tempfile
+
+from ray import tune
+from ray.tune import Checkpoint, RunConfig
+from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
+from ray.tune.tune_config import TuneConfig
+from ray.tune.tuner import Tuner
+
+
+def func(config):
+    starting_epoch = 0
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            with open(os.path.join(checkpoint_dir, "ckpt.pkl"), "rb") as f:
+                checkpoint_dict = pickle.load(f)
+        checkpoint_epoch = checkpoint_dict["epoch"]
+        starting_epoch = checkpoint_epoch + 1
+
+    for epoch in range(starting_epoch, 1000):
+        checkpoint_dict = {"epoch": epoch, "large_data": np.zeros(10000000)}
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with open(os.path.join(tmpdir, "ckpt.pkl"), "wb") as f:
+                pickle.dump(checkpoint_dict, f)
+            tune.report({}, checkpoint=Checkpoint.from_directory(tmpdir))
+
+
+class FrequentPausesScheduler(FIFOScheduler):
+    def on_trial_result(self, tune_controller, trial, result):
+        return TrialScheduler.PAUSE
+
+
+tuner = Tuner(
+    func,
+    tune_config=TuneConfig(num_samples=2, scheduler=FrequentPausesScheduler()),
+    run_config=RunConfig(storage_path="/mnt/cluster_storage", name="frequent_pausing"),
+)
+
+tuner.fit()