chore: import upstream snapshot with attribution

python/ray/tune/examples/custom_func_checkpointing.py

@@ -0,0 +1,70 @@
+# If want to use checkpointing with a custom training function (not a Ray
+# integration like PyTorch or Tensorflow), your function can read/write
+# checkpoint through the ``ray.tune.report(metrics, checkpoint=...)`` API.
+import argparse
+import json
+import os
+import tempfile
+import time
+
+from ray import tune
+from ray.tune import Checkpoint
+
+
+def evaluation_fn(step, width, height):
+    time.sleep(0.1)
+    return (0.1 + width * step / 100) ** (-1) + height * 0.1
+
+
+def train_func(config):
+    step = 0
+    width, height = config["width"], config["height"]
+
+    checkpoint = tune.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as checkpoint_dir:
+            with open(os.path.join(checkpoint_dir, "checkpoint.json")) as f:
+                state = json.load(f)
+            step = state["step"] + 1
+
+    for current_step in range(step, 100):
+        intermediate_score = evaluation_fn(current_step, width, height)
+
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            with open(os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w") as f:
+                json.dump({"step": current_step}, f)
+            tune.report(
+                {"iterations": current_step, "mean_loss": intermediate_score},
+                checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
+            )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing"
+    )
+    args, _ = parser.parse_known_args()
+
+    tuner = tune.Tuner(
+        train_func,
+        run_config=tune.RunConfig(
+            name="hyperband_test",
+            stop={"training_iteration": 1 if args.smoke_test else 10},
+        ),
+        tune_config=tune.TuneConfig(
+            metric="mean_loss",
+            mode="min",
+            num_samples=5,
+        ),
+        param_space={
+            "steps": 10,
+            "width": tune.randint(10, 100),
+            "height": tune.loguniform(10, 100),
+        },
+    )
+    results = tuner.fit()
+    best_result = results.get_best_result()
+    print("Best hyperparameters: ", best_result.config)
+    best_checkpoint = best_result.checkpoint
+    print("Best checkpoint: ", best_checkpoint)