fix: use volume version metadata when updating volumes (#279)

* encode volume content version in source prefix

* use version from metadata

* lint

* format

Author: narasul

src/together/lib/cli/api/beta/jig/_config.py

@@ -6,7 +6,7 @@
 import sys
 import json
 import typing
-from typing import TYPE_CHECKING, Any, Union, Optional
+from typing import TYPE_CHECKING, Any, Union, Optional, cast
 from pathlib import Path
 from dataclasses import field, asdict, dataclass, is_dataclass
 
@@ -60,6 +60,7 @@ class VolumeMount:
 
     name: str
     mount_path: str
+    version: Optional[int] = None
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> VolumeMount:
@@ -84,7 +85,7 @@ class DeployConfig:
     port: int = 8000
     environment_variables: dict[str, str] = field(default_factory=dict[str, str])
     command: Optional[list[str]] = None
-    autoscaling: dict[str, str] = field(default_factory=dict[str, str])
+    autoscaling: dict[str, Union[str, float]] = field(default_factory=dict[str, Union[str, float]])
     health_check_path: str = "/health"
     termination_grace_period_seconds: int = 300
     volume_mounts: list[VolumeMount] = field(default_factory=list[VolumeMount])
@@ -93,7 +94,7 @@ class DeployConfig:
     def from_dict(cls, data: dict[str, Any]) -> DeployConfig:
         deploy_config = {k: v for k, v in data.items() if k in cls.__annotations__}
         if isinstance((mounts := deploy_config.get("volume_mounts")), list):
-            deploy_config["volume_mounts"] = [VolumeMount.from_dict(vm) for vm in mounts]  # pyright: ignore
+            deploy_config["volume_mounts"] = [VolumeMount.from_dict(vm) for vm in cast(list[dict[str, Any]], mounts)]
         return cls(**deploy_config)
 
 
@@ -133,7 +134,7 @@ def validate(value: Any, value_type: type, path: str = "") -> str | None:
         return None
 
     if not isinstance(value, value_type):
-        return f"{path}: expected {type(value).__name__}, got {value!r}"
+        return f"{path}: expected {value_type.__name__}, got {value!r}"  # type: ignore[union-attr]
     return None
 
 
@@ -196,7 +197,7 @@ def load(cls, data: dict[str, Any], path: Path) -> Config:
         else:
             jig_config = data
             if name := jig_config.get("name"):
-                tip = "update `name` in {path}"
+                tip = f"update `name` in {path}"
             else:
                 name = path.resolve().parent.name
                 tip = f"rename your folder or add `name` to {path}"

src/together/lib/cli/api/beta/jig/_utils.py

@@ -89,7 +89,7 @@ def format_deployment_status(d: Deployment) -> str:
         events_status = "\nReplica Events:\n"
         images = set(map(lambda x: x.image or "-", d.replica_events.values()))
         for image in reversed(sorted(images)):
-            events = filter(lambda x: ((x[1].image or "-") == image), d.replica_events.items())
+            events = filter(lambda x: (x[1].image or "-") == image, d.replica_events.items())
             events_status += f"{_image_tag(image)}:\n"
             for replica_id, event in events:
                 events_status += f"  {replica_id}: "

src/together/lib/cli/api/beta/jig/volumes.py

@@ -6,7 +6,7 @@
 import time
 import asyncio
 import itertools
-from typing import Any
+from typing import Any, cast
 from pathlib import Path
 
 import click
@@ -107,13 +107,13 @@ async def spinner_updater(self) -> None:
                 self.update_progress()
             await asyncio.sleep(0.1)
 
-    async def upload_files(self, source_path: Path, volume_name: str) -> None:
+    async def upload_files(self, source_path: Path, volume_name: str, content_version: int) -> None:
         """Upload all files from source directory with progress tracking"""
         # these require a running event loop
         self.semaphore = asyncio.Semaphore(UPLOAD_CONCURRENCY_LIMIT)
         self.progress_lock = asyncio.Lock()
 
-        source_prefix = f"{volume_name}/{source_path.name}"
+        source_prefix = f"{volume_name}/{content_version}"
         files_to_upload: list[tuple[Path, str, int]] = []
 
         for file_path in source_path.rglob("*"):
@@ -296,7 +296,8 @@ async def _create_volume(client: Together, name: str, source: str) -> None:
     if not source_path.is_dir():
         raise ValueError(f"Source path must be a directory: {source}")
 
-    source_prefix = f"{name}/{source_path.name}"
+    content_version = 0
+    source_prefix = f"{name}/{content_version}"
 
     click.echo(f"\N{ROCKET} Creating volume '{name}' with source prefix '{source_prefix}'")
     try:
@@ -310,7 +311,7 @@ async def _create_volume(client: Together, name: str, source: str) -> None:
         raise RuntimeError(f"Failed to create volume: {e}") from e
 
     try:
-        await Uploader(client).upload_files(source_path, volume_name=name)
+        await Uploader(client).upload_files(source_path, volume_name=name, content_version=content_version)
     except Exception as e:
         click.echo(f"\N{CROSS MARK} Upload failed: {e}")
         click.echo(f"\N{WASTEBASKET} Cleaning up volume '{name}'")
@@ -330,16 +331,20 @@ async def _update_volume(client: Together, name: str, source: str) -> None:
         raise ValueError(f"Source path must be a directory: {source}")
 
     try:
-        client.beta.jig.volumes.retrieve(name)
+        response = client.beta.jig.volumes.with_raw_response.retrieve(name)
+        volume_data = response.json()
     except APIStatusError as e:
         if hasattr(e, "status_code") and e.status_code == 404:
             raise ValueError(f"Volume '{name}' does not exist") from e
         raise
 
-    source_prefix = f"{name}/{source_path.name}"
+    volume_dict = cast(dict[str, Any], volume_data)
+    current_version: int = int(volume_dict.get("current_version", 0))
+    version: int = current_version + 1
+    source_prefix = f"{name}/{version}"
 
     click.echo(f"\N{INFORMATION SOURCE} Uploading files for volume '{name}'")
-    await Uploader(client).upload_files(source_path, volume_name=name)
+    await Uploader(client).upload_files(source_path, volume_name=name, content_version=version)
 
     click.echo(f"\N{INFORMATION SOURCE} Updating volume '{name}' with source prefix '{source_prefix}'")
     client.beta.jig.volumes.update(