chore: Update client.py & backend.py

Author: bradhilton

src/art/client.py

@@ -3,7 +3,6 @@
 from typing import AsyncIterator, Literal, TypedDict, cast
 
 import httpx
-from openai import AsyncOpenAI, BaseModel, _exceptions
 from openai._base_client import AsyncAPIClient, AsyncPaginator, make_request_options
 from openai._compat import cached_property
 from openai._qs import Querystring
@@ -16,19 +15,27 @@
 from openai.resources.models import AsyncModels  # noqa: F401
 from typing_extensions import override
 
+from openai import AsyncOpenAI, BaseModel, _exceptions
+
 from .trajectories import TrajectoryGroup
 
 
 class Checkpoint(BaseModel):
     id: str
     model_id: str
     step: int
+    metrics: dict[str, float]
 
 
 class CheckpointListParams(TypedDict, total=False):
     model_id: str
 
 
+class DeleteCheckpointsResponse(BaseModel):
+    deleted_count: int
+    not_found_steps: list[int]
+
+
 class Checkpoints(AsyncAPIResource):
     async def retrieve(
         self, *, model_id: str, step: int | Literal["latest"]
@@ -64,6 +71,16 @@ def list(
             model=Checkpoint,
         )
 
+    async def delete(
+        self, *, model_id: str, steps: list[int]
+    ) -> DeleteCheckpointsResponse:
+        return await self._delete(
+            f"/preview/models/{model_id}/checkpoints",
+            body={"steps": steps},
+            cast_to=DeleteCheckpointsResponse,
+            options=dict(max_retries=0),
+        )
+
 
 class Model(BaseModel):
     id: str
@@ -128,6 +145,7 @@ async def create(
                     "base_model": base_model,
                     "return_existing": return_existing,
                 },
+                options=dict(max_retries=0),
             )
         )
 
@@ -229,6 +247,7 @@ async def create(
                 ],
                 "experimental_config": experimental_config,
             },
+            options=dict(max_retries=0),
         )
 
     async def retrieve(self, training_job_id: int) -> TrainingJob:

src/art/serverless/backend.py

@@ -69,7 +69,22 @@ async def _delete_checkpoints(
         benchmark: str,
         benchmark_smoothing: float,
     ) -> None:
-        raise NotImplementedError
+        # TODO: potentially implement benchmark smoothing
+        max_metric: float | None = None
+        max_step: int | None = None
+        all_steps: list[int] = []
+        async for checkpoint in self._client.checkpoints.list(model_id=model.id):
+            metric = checkpoint.metrics.get(benchmark, None)
+            if metric is not None and (max_metric is None or metric > max_metric):
+                max_metric = metric
+                max_step = checkpoint.step
+            all_steps.append(checkpoint.step)
+        steps_to_delete = [step for step in all_steps[:-1] if step != max_step]
+        if steps_to_delete:
+            await self._client.checkpoints.delete(
+                model_id=model.id,
+                steps=steps_to_delete,
+            )
 
     async def _prepare_backend_for_training(
         self,