fix: Refactor temporal clue example to enhance S3 backup handling and streamline training loop

- Updated the `temporal-clue.py` script to check for the `BACKUP_BUCKET` environment variable before pulling from and pushing to S3, improving robustness.
- Simplified the training loop structure for better readability and maintainability.
- Adjusted the `LocalBackend` class to handle potential import errors gracefully when deleting checkpoints.

This change enhances the overall functionality and error handling of the training process.

Author: bradhilton

examples/temporal_clue/temporal-clue.py

@@ -60,35 +60,47 @@ async def main():
         base_model="Qwen/Qwen2.5-7B-Instruct",
         _internal_config={"init_args": {"gpu_memory_utilization": 0.775}},
     )
-    backend = LocalBackend()
-    await backend._experimental_pull_from_s3(model)
-    await model.register(backend)
+    with LocalBackend() as backend:
+        if "BACKUP_BUCKET" in os.environ:
+            await backend._experimental_pull_from_s3(
+                model,
+                s3_bucket=os.environ["BACKUP_BUCKET"],
+                verbose=True,
+            )
+        else:
+            print("BACKUP_BUCKET not found in environment variables")
+        await model.register(backend)
 
-    stride = 4
-    for i in range(await model.get_step(), 1_000):
-        val_groups, train_groups = await asyncio.gather(
-            art.gather_trajectory_groups(
-                (
-                    art.TrajectoryGroup(rollout(model, puzzle) for _ in range(2))
-                    for puzzle in val_puzzles
+        stride = 4
+        for i in range(await model.get_step(), 1_000):
+            val_groups, train_groups = await asyncio.gather(
+                art.gather_trajectory_groups(
+                    (
+                        art.TrajectoryGroup(rollout(model, puzzle) for _ in range(2))
+                        for puzzle in val_puzzles
+                    ),
+                    pbar_desc="val",
                 ),
-                pbar_desc="val",
-            ),
-            art.gather_trajectory_groups(
-                (
-                    art.TrajectoryGroup(rollout(model, puzzle) for _ in range(50))
-                    for puzzle in train_puzzles[i * stride : (i + 1) * stride]
+                art.gather_trajectory_groups(
+                    (
+                        art.TrajectoryGroup(rollout(model, puzzle) for _ in range(50))
+                        for puzzle in train_puzzles[i * stride : (i + 1) * stride]
+                    ),
+                    pbar_desc="train",
                 ),
-                pbar_desc="train",
-            ),
-        )
-        await model.log(val_groups)
-        await model.delete_checkpoints()
-        await backend._experimental_push_to_s3(model)
-        result = await backend.train(model, train_groups, learning_rate=5e-5)
-        await model.log(
-            train_groups, metrics=result.metrics, step=result.step, split="train"
-        )
+            )
+            await model.log(val_groups)
+            await model.delete_checkpoints()
+            if "BACKUP_BUCKET" in os.environ:
+                await backend._experimental_push_to_s3(
+                    model,
+                    s3_bucket=os.environ["BACKUP_BUCKET"],
+                    verbose=True,
+                )
+            result = await backend.train(model, train_groups, learning_rate=5e-5)
+            await model.log(
+                train_groups, metrics=result.metrics, step=result.step, split="train"
+            )
 
 
 if __name__ == "__main__":

src/art/local/backend.py

@@ -148,7 +148,7 @@ def _model_inference_name(self, model: Model, step: int | None = None) -> str:
         # Default to step 0 when not specified (the initial checkpoint created at registration)
         if step is not None:
             actual_step = step
-        elif model.name in self._services:
+        elif model.name in self._services and self._in_process:
             # In dedicated mode the service tracks which adapter vLLM has
             # actually loaded.  Reading the filesystem would race: the
             # checkpoint directory appears before the HTTP reload completes.
@@ -294,14 +294,18 @@ async def _delete_checkpoint_files(
         steps_to_keep: list[int],
     ) -> None:
         """Delete checkpoint files, keeping only the specified steps."""
-        from ..tinker.service import TinkerService
 
         output_dir = get_model_dir(model=model, art_path=self._path)
         service = await self._get_service(model)
-        if isinstance(service, TinkerService):
-            await service.delete_checkpoints(steps_to_keep)
-        else:
-            delete_checkpoints(output_dir, steps_to_keep)
+        try:
+            from ..tinker.service import TinkerService
+
+            if isinstance(service, TinkerService):
+                await service.delete_checkpoints(steps_to_keep)
+                return
+        except ImportError:
+            pass
+        delete_checkpoints(output_dir, steps_to_keep)
 
     async def _prepare_backend_for_training(
         self,