Update Checkpointing Options construction through CLI configuration to move from Orbax V0 CheckpointManagerOptions to Tunix's CheckpointOptions.

docs/launching.md

@@ -344,6 +344,8 @@ before performing a parameter update (simulates larger batch sizes).
 * **`checkpointing_options`**:
   * `max_to_keep`: Number of recent checkpoints to retain.
   * `save_interval_steps`: How often to save a checkpoint.
+  * `enable_async_checkpointing`: Boolean to toggle asynchronous checkpointing execution.
+  * `timeout_secs`: Maximum time permitted for asynchronous writes natively.
 
 
 * **`metrics_logging_options`**: Settings for logging. Includes project name, run name, and flush frequency.

docs/reliability.md

@@ -21,6 +21,31 @@ training step count. By default, checkpointing is disabled if
 `checkpoint_root_directory` is not specified. Users can further customize
 checkpointing behavior via `checkpointing_options` in the config.
 
+Users customize background preservation behavior granularly using components
+defined inside `checkpoint_options`:
+
+*   **Save Decision Policies**: Dictates when to initiate a checkpoint based on
+    defined steps or intervals. Supported configurations include
+    `FixedIntervalPolicy` and `ContinuousCheckpointingPolicy`. The default is
+    `ContinuousCheckpointingPolicy(minimum_interval_secs=180)` (saves every 180
+    seconds). See Orbax v1 [`save_decision_policies.py`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/experimental/v1/_src/training/save_decision_policies.py)
+    for the complete interface contracts.
+*   **Preservation Policies**: Sets specifications regarding tracking
+    checkpoints over bounded timelines (e.g., `LatestN`). The default is
+    `LatestN(n=3)` (keeps the latest 3 checkpoints). See Orbax v1
+    [`preservation_policies.py`](https://github.com/google/orbax/blob/main/checkpoint/orbax/checkpoint/experimental/v1/_src/training/preservation_policies.py)
+    for the complete interface contracts.
+*   **Step Name Format**: Defines the representation of directory names for step
+    checkpoints. The default is `ocp.path.step.standard_name_format()` (uses
+    simple integer step names).
+*   **Asynchronous Processing**: Manage asynchronous behavior by specifying:
+    *   `enable_async_checkpointing`: Whether to use async checkpointing.
+        Defaults to `True`. **It is recommended to keep this enabled** to
+        prevent the main thread from blocking during training runs while
+        checkpoints are written to storage.
+    *   `timeout_secs`: The timeout for asynchronous operations.
+        Defaults to `1200` seconds.
+
 ## Fault Tolerance
 
 Tunix ensures fault tolerance primarily through its checkpointing mechanism,

tests/cli/config_test.py

@@ -699,6 +699,52 @@ def test_dict_to_cli_args_with_none(self):
     got = list(config._dict_to_cli_args(d))
     self.assertEqual(expected, got)
 
+  def test_obtain_training_config_dict_checkpointing(self):
+    hp = self.initialize_config([])
+
+    # Valid options
+    hp.config["training_config"] = {
+        "checkpointing_options": {
+            "save_interval_steps": 10,
+            "max_to_keep": 5,
+            "enable_async_checkpointing": True,
+            "timeout_secs": 60,
+        }
+    }
+    result = hp.obtain_training_config_dict("training_config")
+    self.assertIn("checkpointing_options", result)
+    opts = result["checkpointing_options"]
+    self.assertEqual(opts.enable_async_checkpointing, True)
+    self.assertEqual(opts.async_options.timeout_secs, 60)
+    self.assertIsNotNone(opts.save_decision_policy)
+    self.assertIsNotNone(opts.preservation_policy)
+
+  def test_obtain_training_config_dict_checkpointing_invalid_options(self):
+    hp = self.initialize_config([])
+    hp.config["training_config"] = {
+        "checkpointing_options": "not a dict"
+    }
+    with self.assertRaisesRegex(
+        ValueError, "Expected dictionary for checkpointing_options"
+    ):
+      hp.obtain_training_config_dict("training_config")
+
+    hp.config["training_config"] = {
+        "checkpointing_options": {
+            "invalid_key": 10,
+        }
+    }
+    with self.assertRaisesRegex(ValueError, "Invalid checkpointing options"):
+      hp.obtain_training_config_dict("training_config")
+
+    hp.config["training_config"] = {
+        "checkpointing_options": ["save_interval_steps"]
+    }
+    with self.assertRaisesRegex(
+        ValueError, "Expected dictionary for checkpointing_options"
+    ):
+      hp.obtain_training_config_dict("training_config")
+
 
 if __name__ == "__main__":
   if "HF_TOKEN" not in os.environ: