chore: import upstream snapshot with attribution

doc/source/train/doc_code/key_concepts.py

@@ -0,0 +1,136 @@
+# flake8: noqa
+# isort: skip_file
+
+from pathlib import Path
+import tempfile
+
+import ray.train
+from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
+
+
+def train_fn(config):
+    for i in range(3):
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            Path(temp_checkpoint_dir).joinpath("model.pt").touch()
+            ray.train.report(
+                {"loss": i},
+                checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
+            )
+
+    return {"total loss": 3}
+
+
+trainer = DataParallelTrainer(
+    train_fn, scaling_config=ray.train.ScalingConfig(num_workers=2)
+)
+
+
+# __run_config_start__
+import os
+
+from ray.train import RunConfig
+
+run_config = RunConfig(
+    # Name of the training run (directory name).
+    name="my_train_run",
+    # The experiment results will be saved to: storage_path/name
+    storage_path=os.path.expanduser("~/ray_results"),
+    # storage_path="s3://my_bucket/tune_results",
+)
+# __run_config_end__
+
+# __checkpoint_config_start__
+from ray.train import RunConfig, CheckpointConfig
+
+# Example 1: Only keep the 2 *most recent* checkpoints and delete the others.
+run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=2))
+
+
+# Example 2: Only keep the 2 *best* checkpoints and delete the others.
+run_config = RunConfig(
+    checkpoint_config=CheckpointConfig(
+        num_to_keep=2,
+        # *Best* checkpoints are determined by these params:
+        checkpoint_score_attribute="mean_accuracy",
+        checkpoint_score_order="max",
+    ),
+    # This will store checkpoints on S3.
+    storage_path="s3://remote-bucket/location",
+)
+# __checkpoint_config_end__
+
+
+# __result_metrics_start__
+result = trainer.fit()
+
+print("Observed metrics:", result.metrics)
+# __result_metrics_end__
+
+
+# __result_dataframe_start__
+df = result.metrics_dataframe
+print("Minimum loss", min(df["loss"]))
+# __result_dataframe_end__
+
+
+# __result_return_value_start__
+print("Returned data", result.return_value)
+# __result_return_value_end__
+
+
+# __result_checkpoint_start__
+print("Last checkpoint:", result.checkpoint)
+
+with result.checkpoint.as_directory() as tmpdir:
+    # Load model from directory
+    ...
+# __result_checkpoint_end__
+
+# __result_best_checkpoint_start__
+# Print available checkpoints
+for checkpoint, metrics in result.best_checkpoints:
+    print("Loss", metrics["loss"], "checkpoint", checkpoint)
+
+# Get checkpoint with minimal loss
+best_checkpoint = min(
+    result.best_checkpoints, key=lambda checkpoint: checkpoint[1]["loss"]
+)[0]
+
+with best_checkpoint.as_directory() as tmpdir:
+    # Load model from directory
+    ...
+# __result_best_checkpoint_end__
+
+import pyarrow
+
+# __result_path_start__
+result_path: str = result.path
+result_filesystem: pyarrow.fs.FileSystem = result.filesystem
+
+print(f"Results location (fs, path) = ({result_filesystem}, {result_path})")
+# __result_path_end__
+
+
+# __result_restore_start__
+from ray.train import Result
+
+restored_result = Result.from_path(result_path)
+print("Restored loss", restored_result.metrics["loss"])
+# __result_restore_end__
+
+
+def error_train_fn(config):
+    raise RuntimeError("Simulated training error")
+
+
+trainer = DataParallelTrainer(
+    error_train_fn, scaling_config=ray.train.ScalingConfig(num_workers=1)
+)
+
+# __result_error_start__
+try:
+    result = trainer.fit()
+except ray.train.TrainingFailedError as e:
+    if isinstance(e, ray.train.WorkerGroupError):
+        print(e.worker_failures)
+# __result_error_end__