Fix remaining issues with tuning

climanet/predict.py

@@ -59,6 +59,7 @@ def predict_monthly_var(
     run_dir: str = ".",
     verbose: bool = True,
     dataloader_num_workers: int = 2,
+    store_logs: bool = True,
 ):
     """
     Predicts monthly variable values using a trained model and a provided dataset.
@@ -105,7 +106,8 @@ def predict_monthly_var(
     all_predictions = torch.empty(len(dataset), M, H, W, device=device)
 
     # Set up logging
-    writer = setup_logging(run_dir)
+    if store_logs:
+        writer = setup_logging(run_dir)
 
     with torch.inference_mode():
         idx = 0
@@ -140,13 +142,16 @@ def predict_monthly_var(
                     f"Processed batch {i + 1}/{len(dataloader)}, with loss: {loss.item():.4f}"
                 )
 
-            writer.add_scalar("Progress/Batch", i + 1, idx)
+            if store_logs:
+                writer.add_scalar("Progress/Batch", i + 1, idx)
 
     average_loss = average_loss.item() / len(dataloader)
 
     if verbose:
         print(f"Average loss over all batches: {average_loss:.4f}")
-    writer.add_scalar("Loss/Average", average_loss)
+
+    if store_logs:
+        writer.add_scalar("Loss/Average", average_loss)
 
     if return_numpy:
         all_predictions = all_predictions.cpu().numpy()
@@ -159,10 +164,12 @@ def predict_monthly_var(
         if verbose:
             print(f"Predictions saved to '{run_dir}'")
 
-        writer.add_text("Info", f"Predictions saved to '{run_dir}'")
+        if store_logs:
+            writer.add_text("Info", f"Predictions saved to '{run_dir}'")
 
     # Close the writer when done
-    writer.close()
+    if store_logs:
+        writer.close()
 
     if return_loss:
         all_predictions = (all_predictions, average_loss)

climanet/train.py

@@ -1,3 +1,4 @@
+import tempfile
 from pathlib import Path
 
 import torch
@@ -41,6 +42,7 @@ def train_monthly_model(
     dataloader_num_workers: int = 2,
     verbose_epoch_interval: int = 20,
     tune_checkpoint: bool = False,
+    store_logs: bool = True,
 ):
     """Train the model to predict monthly data from daily data.
     Args:
@@ -75,7 +77,8 @@ def train_monthly_model(
     )
 
     # Set up logging
-    writer = setup_logging(run_dir)
+    if store_logs:
+        writer = setup_logging(run_dir)
 
     # Set the optimizer
     optimizer = torch.optim.AdamW(
@@ -133,7 +136,10 @@ def train_monthly_model(
 
         # Calculate average epoch loss
         avg_train_loss = epoch_loss.item() / (i + 1)
-        writer.add_scalar("Loss/train", avg_train_loss, epoch)
+
+        if store_logs:
+            writer.add_scalar("Loss/train", avg_train_loss, epoch)
+
         avg_epoch_loss = avg_train_loss  # Initially use training loss
 
         # Validation loss (optional)
@@ -150,10 +156,13 @@ def train_monthly_model(
                 verbose=False,
                 run_dir=run_dir,
                 dataloader_num_workers=dataloader_num_workers,
+                store_logs=False,
             )
-            writer.add_scalar("Loss/validation", avg_val_loss, epoch)
             avg_epoch_loss = avg_val_loss  # Use validation loss if exists
 
+            if store_logs:
+                writer.add_scalar("Loss/validation", avg_val_loss, epoch)
+
             if verbose and epoch % verbose_epoch_interval == 0:
                 gap = avg_val_loss - avg_train_loss
                 print(f"Epoch {epoch}: gap between train and val loss: {gap:.6f}")
@@ -173,30 +182,36 @@ def train_monthly_model(
             counter += 1
 
         # Log to TensorBoard
-        writer.add_scalar("Loss/best", best_loss, epoch)
+        if store_logs:
+            writer.add_scalar("Loss/best", best_loss, epoch)
 
         if verbose and epoch % 20 == 0:
             print(f"Epoch {epoch}: best_loss = {best_loss:.6f}")
 
         # Only stop if LR is at minimum AND no improvement
         current_lr = optimizer.param_groups[0]["lr"]
         if counter >= patience and current_lr <= scheduler.min_lrs[0]:
-            writer.add_text("Training", f"Early stop at epoch {epoch}", epoch)
+            if store_logs:
+                writer.add_text("Training", f"Early stop at epoch {epoch}", epoch)
             break
 
     # Restore best model
     if best_state_dict is not None:
         model.load_state_dict(best_state_dict)
 
     if tune_checkpoint:
-        # Save the model and optimizer state for Ray Tune checkpointing
-        save_model(model, optimizer, run_dir, filename="checkpoint.pt", verbose=False)
-        tune.report(
-            {"loss": best_loss}, checkpoint=tune.Checkpoint.from_directory(run_dir)
-        )
+        with tempfile.TemporaryDirectory() as checkpoint_dir:
+            checkpoint_path = Path(checkpoint_dir)
+
+            # Save the model and optimizer state for Ray Tune checkpointing
+            save_model(model, optimizer, checkpoint_path, filename="checkpoint.pt", verbose=False)
+            tune.report(
+                {"loss": best_loss}, checkpoint=tune.Checkpoint.from_directory(checkpoint_path)
+            )
 
     # Close the writer when done
-    writer.close()
+    if store_logs:
+        writer.close()
 
     if verbose:
         print(f"Training complete. Best loss: {best_loss:.6f}")

climanet/tune.py

@@ -2,6 +2,7 @@
 
 import ray
 import xarray as xr
+from ray.air.config import CheckpointConfig
 from ray.tune.schedulers import ASHAScheduler
 
 from climanet.dataset import STDataset
@@ -20,8 +21,15 @@ def _train(tune_config, static_args):
     num_epoch = static_args["num_epoch"]
 
     # dont use ray.put() and ray.get() (i.e. object store) when data is large
-    train_dataset = tune_data_preparation(static_args["data_config_train"])
-    validation_dataset = tune_data_preparation(static_args["data_config_validation"])
+    if static_args.get("train_dataset") is not None:
+        train_dataset = ray.get(static_args["train_dataset"])
+    else:
+        train_dataset = tune_data_preparation(static_args["data_config_train"])
+
+    if static_args.get("validation_dataset") is not None:
+        validation_dataset = ray.get(static_args["validation_dataset"])
+    else:
+        validation_dataset = tune_data_preparation(static_args["data_config_validation"])
 
     patch_size = tune_config["patch_size"]
     overlap = tune_config["overlap"]
@@ -61,6 +69,7 @@ def _train(tune_config, static_args):
         store_model=False,
         verbose=False,
         tune_checkpoint=True,
+        store_logs=False,
     )
 
 
@@ -117,6 +126,13 @@ def run_tune(tune_config: dict, static_args: dict):
     if Path(experiment_path).exists():
         tuner = ray.tune.Tuner.restore(
             experiment_path,
+            ray.tune.with_resources(
+                ray.tune.with_parameters(_train, static_args=static_args),
+                resources={
+                    "cpu": static_args["cpu_per_trial"],
+                    "gpu": static_args["gpu_per_trial"],
+                },
+            ),
             resume_errored=True,
         )
     else:
@@ -136,7 +152,15 @@ def run_tune(tune_config: dict, static_args: dict):
                 max_concurrent_trials=static_args["max_concurrent_trials"],
             ),
             param_space=tune_config,
-            run_config=ray.tune.RunConfig(storage_path=static_args["run_dir"], name=experiment_name),
+            run_config=ray.tune.RunConfig(
+                storage_path=static_args["run_dir"],
+                name=experiment_name,
+                checkpoint_config=CheckpointConfig(
+                    num_to_keep=1,
+                    checkpoint_score_attribute="loss",
+                    checkpoint_score_order="min",
+                ),
+            ),
         )
 
     results = tuner.fit()

notebooks/example_tuning.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "id": "57546057-2042-42eb-a793-23e77d22965e",
    "metadata": {},
    "outputs": [],
@@ -168,8 +168,8 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2026-07-23 10:11:28,180\tINFO worker.py:2024 -- Started a local Ray instance.\n",
-      "\u001b[36m(_train pid=40937)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56/_train_2bdd9_00001_1_patch_size=4_2026-07-23_10-11-56/checkpoint_000000)\n"
+      "2026-08-03 10:59:39,452\tINFO worker.py:2024 -- Started a local Ray instance.\n",
+      "\u001b[36m(_train pid=127429)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune/_train_ac0e6_00001_1_patch_size=4_2026-08-03_10-59-44/checkpoint_000000)\n"
      ]
     }
    ],
@@ -185,7 +185,8 @@
     "    \"train_dataset\": ray.put(train_dataset),\n",
     "    \"validation_dataset\": ray.put(validation_dataset),\n",
     "    \"num_epoch\": 1,\n",
-    "    \"max_concurrent_trials\": 2\n",
+    "    \"max_concurrent_trials\": 2,\n",
+    "    \"experiment_name\": \"climanet_tune\",\n",
     "}\n",
     "\n",
     "# parameters to tune\n",
@@ -225,9 +226,9 @@
        "      <h3>Tune Status</h3>\n",
        "      <table>\n",
        "<tbody>\n",
-       "<tr><td>Current time:</td><td>2026-07-23 10:12:08</td></tr>\n",
-       "<tr><td>Running for: </td><td>00:00:11.76        </td></tr>\n",
-       "<tr><td>Memory:      </td><td>8.6/15.3 GiB       </td></tr>\n",
+       "<tr><td>Current time:</td><td>2026-08-03 10:59:57</td></tr>\n",
+       "<tr><td>Running for: </td><td>00:00:13.08        </td></tr>\n",
+       "<tr><td>Memory:      </td><td>9.8/15.3 GiB       </td></tr>\n",
        "</tbody>\n",
        "</table>\n",
        "    </div>\n",
@@ -243,11 +244,11 @@
        "    <h3>Trial Status</h3>\n",
        "    <table>\n",
        "<thead>\n",
-       "<tr><th>Trial name        </th><th>status    </th><th>loc               </th><th style=\"text-align: right;\">  patch_size</th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th><th style=\"text-align: right;\">    loss</th></tr>\n",
+       "<tr><th>Trial name        </th><th>status    </th><th>loc                </th><th style=\"text-align: right;\">  patch_size</th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th><th style=\"text-align: right;\">    loss</th></tr>\n",
        "</thead>\n",
        "<tbody>\n",
-       "<tr><td>_train_2bdd9_00000</td><td>TERMINATED</td><td>192.168.2.13:40938</td><td style=\"text-align: right;\">           2</td><td style=\"text-align: right;\">     1</td><td style=\"text-align: right;\">         6.45845</td><td style=\"text-align: right;\">0.384662</td></tr>\n",
-       "<tr><td>_train_2bdd9_00001</td><td>TERMINATED</td><td>192.168.2.13:40937</td><td style=\"text-align: right;\">           4</td><td style=\"text-align: right;\">     1</td><td style=\"text-align: right;\">         3.89669</td><td style=\"text-align: right;\">0.365478</td></tr>\n",
+       "<tr><td>_train_ac0e6_00000</td><td>TERMINATED</td><td>192.168.2.13:127430</td><td style=\"text-align: right;\">           2</td><td style=\"text-align: right;\">     1</td><td style=\"text-align: right;\">         6.44704</td><td style=\"text-align: right;\">0.384662</td></tr>\n",
+       "<tr><td>_train_ac0e6_00001</td><td>TERMINATED</td><td>192.168.2.13:127429</td><td style=\"text-align: right;\">           4</td><td style=\"text-align: right;\">     1</td><td style=\"text-align: right;\">         3.76473</td><td style=\"text-align: right;\">0.365478</td></tr>\n",
        "</tbody>\n",
        "</table>\n",
        "  </div>\n",
@@ -294,9 +295,15 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2026-07-23 10:12:08,344\tINFO tune.py:1007 -- Wrote the latest version of all result files and experiment state to '/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56' in 0.0033s.\n",
-      "2026-07-23 10:12:08,348\tINFO tune.py:1039 -- Total run time: 11.79 seconds (11.75 seconds for the tuning loop).\n",
-      "\u001b[36m(_train pid=40938)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56/_train_2bdd9_00000_0_patch_size=2_2026-07-23_10-11-56/checkpoint_000000)\n"
+      "2026-08-03 10:59:55,339\tWARNING experiment_state.py:209 -- Experiment state snapshotting has been triggered multiple times in the last 5.0 seconds and may become a bottleneck. A snapshot is forced if `CheckpointConfig(num_to_keep)` is set, and a trial has checkpointed >= `num_to_keep` times since the last snapshot.\n",
+      "You may want to consider increasing the `CheckpointConfig(num_to_keep)` or decreasing the frequency of saving checkpoints.\n",
+      "You can suppress this warning by setting the environment variable TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S to a smaller value than the current threshold (5.0). Set it to 0 to completely suppress this warning.\n",
+      "2026-08-03 10:59:57,963\tWARNING experiment_state.py:209 -- Experiment state snapshotting has been triggered multiple times in the last 5.0 seconds and may become a bottleneck. A snapshot is forced if `CheckpointConfig(num_to_keep)` is set, and a trial has checkpointed >= `num_to_keep` times since the last snapshot.\n",
+      "You may want to consider increasing the `CheckpointConfig(num_to_keep)` or decreasing the frequency of saving checkpoints.\n",
+      "You can suppress this warning by setting the environment variable TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S to a smaller value than the current threshold (5.0). Set it to 0 to completely suppress this warning.\n",
+      "2026-08-03 10:59:57,971\tINFO tune.py:1007 -- Wrote the latest version of all result files and experiment state to '/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune' in 0.0066s.\n",
+      "2026-08-03 10:59:57,986\tINFO tune.py:1039 -- Total run time: 13.11 seconds (13.07 seconds for the tuning loop).\n",
+      "\u001b[36m(_train pid=127430)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune/_train_ac0e6_00000_0_patch_size=2_2026-08-03_10-59-44/checkpoint_000000)\n"
      ]
     }
    ],
@@ -325,9 +332,9 @@
       "text/plain": [
        "Result(\n",
        "  metrics={'loss': 0.36547836661338806},\n",
-       "  path='/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56/_train_2bdd9_00001_1_patch_size=4_2026-07-23_10-11-56',\n",
+       "  path='/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune/_train_ac0e6_00001_1_patch_size=4_2026-08-03_10-59-44',\n",
        "  filesystem='local',\n",
-       "  checkpoint=Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-23_10-11-56/_train_2bdd9_00001_1_patch_size=4_2026-07-23_10-11-56/checkpoint_000000)\n",
+       "  checkpoint=Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/climanet_tune/_train_ac0e6_00001_1_patch_size=4_2026-08-03_10-59-44/checkpoint_000000)\n",
        ")"
       ]
      },
@@ -350,10 +357,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 10,
    "id": "66645a30-89d7-459b-9fcc-e03fa4d34b71",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2026-08-03 11:01:08,865\tINFO worker.py:2024 -- Started a local Ray instance.\n"
+     ]
+    }
+   ],
    "source": [
     "# find the path to best model\n",
     "if not ray.is_initialized():\n",
@@ -368,7 +383,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 11,
    "id": "67274359-2e8a-485f-a56a-a3831bd5f9cb",
    "metadata": {},
    "outputs": [],
@@ -394,7 +409,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 12,
    "id": "65c8bab8-dc84-46ee-a249-24b7dc24f991",
    "metadata": {},
    "outputs": [
@@ -404,7 +419,7 @@
        "0.3733350435892741"
       ]
      },
-     "execution_count": 20,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }