K-Clip Upload

kscale/cli.py

@@ -4,6 +4,7 @@
 
 from kscale.utils.cli import recursive_help
 from kscale.web.kernel_images import cli as kernel_images_cli
+from kscale.web.krec import cli as krec_cli
 from kscale.web.pybullet import cli as pybullet_cli
 from kscale.web.urdf import cli as urdf_cli
 
@@ -17,6 +18,7 @@ def cli() -> None:
 cli.add_command(urdf_cli, "urdf")
 cli.add_command(pybullet_cli, "pybullet")
 cli.add_command(kernel_images_cli, "kernel")
+cli.add_command(krec_cli, "krec")
 
 if __name__ == "__main__":
     # python -m kscale.cli

kscale/web/gen/api.py

@@ -2,7 +2,7 @@
 
 # generated by datamodel-codegen:
 #   filename:  openapi.json
-#   timestamp: 2024-11-21T00:20:17+00:00
+#   timestamp: 2024-11-21T21:35:07+00:00
 
 from __future__ import annotations
 
@@ -63,6 +63,10 @@ class ClientIdResponse(BaseModel):
     client_id: str = Field(..., title="Client Id")
 
 
+class CompletedKRecUploadResponse(BaseModel):
+    status: str = Field(..., title="Status")
+
+
 class CreateCheckoutSessionRequest(BaseModel):
     listing_id: str = Field(..., title="Listing Id")
     stripe_product_id: str = Field(..., title="Stripe Product Id")
@@ -130,6 +134,34 @@ class GoogleLogin(BaseModel):
     token: str = Field(..., title="Token")
 
 
+class UploadStatus(Enum):
+    pending = "pending"
+    completed = "completed"
+
+
+class KRec(BaseModel):
+    id: str = Field(..., title="Id")
+    user_id: str = Field(..., title="User Id")
+    robot_id: str = Field(..., title="Robot Id")
+    created_at: int = Field(..., title="Created At")
+    name: str = Field(..., title="Name")
+    description: Optional[str] = Field(None, title="Description")
+    upload_status: Optional[UploadStatus] = Field("pending", title="Upload Status")
+
+
+class KRecPartCompleted(BaseModel):
+    part_number: int = Field(..., title="Part Number")
+    etag: str = Field(..., title="Etag")
+    checksum: Optional[str] = Field(None, title="Checksum")
+
+
+class KRecUrls(BaseModel):
+    url: str = Field(..., title="Url")
+    filename: str = Field(..., title="Filename")
+    expires_at: int = Field(..., title="Expires At")
+    checksum: Optional[str] = Field(None, title="Checksum")
+
+
 class Permission(Enum):
     read = "read"
     write = "write"
@@ -184,6 +216,15 @@ class LoginResponse(BaseModel):
     token: str = Field(..., title="Token")
 
 
+class MultipartUploadDetails(BaseModel):
+    upload_id: str = Field(..., title="Upload Id")
+    presigned_urls: List[Dict[str, Union[str, int]]] = Field(..., title="Presigned Urls")
+    bucket: str = Field(..., title="Bucket")
+    key: str = Field(..., title="Key")
+    part_size: int = Field(..., title="Part Size")
+    num_parts: int = Field(..., title="Num Parts")
+
+
 class Permission1(Enum):
     is_admin = "is_admin"
     is_mod = "is_mod"
@@ -254,6 +295,11 @@ class Order(BaseModel):
     refunded_date: Optional[int] = Field(None, title="Refunded Date")
 
 
+class PresignedUrlResponse(BaseModel):
+    upload_url: str = Field(..., title="Upload Url")
+    artifact_id: str = Field(..., title="Artifact Id")
+
+
 class ProcessPreorderResponse(BaseModel):
     status: str = Field(..., title="Status")
     checkout_session: Dict[str, Any] = Field(..., title="Checkout Session")
@@ -324,18 +370,26 @@ class SetRequest(BaseModel):
 
 
 class ArtifactType(Enum):
+    image = "image"
+
+
+class ArtifactType1(Enum):
+    kernel = "kernel"
+
+
+class ArtifactType2(Enum):
     urdf = "urdf"
     mjcf = "mjcf"
 
 
-class ArtifactType1(Enum):
+class ArtifactType3(Enum):
     stl = "stl"
     obj = "obj"
     dae = "dae"
     ply = "ply"
 
 
-class ArtifactType2(Enum):
+class ArtifactType4(Enum):
     tgz = "tgz"
     zip = "zip"
 
@@ -346,12 +400,27 @@ class SingleArtifactResponse(BaseModel):
     username: str = Field(..., title="Username")
     slug: str = Field(..., title="Slug")
     name: str = Field(..., title="Name")
-    artifact_type: Union[str, ArtifactType, ArtifactType1, ArtifactType2] = Field(..., title="Artifact Type")
+    artifact_type: Union[ArtifactType, ArtifactType1, ArtifactType2, ArtifactType3, ArtifactType4] = Field(
+        ..., title="Artifact Type"
+    )
     description: Optional[str] = Field(..., title="Description")
     timestamp: int = Field(..., title="Timestamp")
     urls: ArtifactUrls
     is_main: Optional[bool] = Field(False, title="Is Main")
     can_edit: Optional[bool] = Field(False, title="Can Edit")
+    size: Optional[int] = Field(None, title="Size")
+
+
+class SingleKRecResponse(BaseModel):
+    id: str = Field(..., title="Id")
+    name: str = Field(..., title="Name")
+    created_at: int = Field(..., title="Created At")
+    user_id: str = Field(..., title="User Id")
+    robot_id: str = Field(..., title="Robot Id")
+    type: Optional[str] = Field("KRec", title="Type")
+    upload_status: str = Field(..., title="Upload Status")
+    urls: Optional[KRecUrls] = None
+    size: Optional[int] = Field(None, title="Size")
 
 
 class SingleRobotResponse(BaseModel):
@@ -434,6 +503,19 @@ class UploadArtifactResponse(BaseModel):
     artifacts: List[SingleArtifactResponse] = Field(..., title="Artifacts")
 
 
+class UploadKRecRequest(BaseModel):
+    name: str = Field(..., title="Name")
+    robot_id: str = Field(..., title="Robot Id")
+    description: Optional[str] = Field(None, title="Description")
+    file_size: Optional[int] = Field(None, title="File Size")
+    part_size: Optional[int] = Field(None, title="Part Size")
+
+
+class UploadKRecResponse(BaseModel):
+    krec_id: str = Field(..., title="Krec Id")
+    upload_details: MultipartUploadDetails
+
+
 class UserInfoResponseItem(BaseModel):
     id: str = Field(..., title="Id")
     email: str = Field(..., title="Email")
@@ -456,6 +538,12 @@ class ValidationError(BaseModel):
     type: str = Field(..., title="Error Type")
 
 
+class CompletedKRecUploadRequest(BaseModel):
+    krec_id: str = Field(..., title="Krec Id")
+    upload_id: str = Field(..., title="Upload Id")
+    parts: List[KRecPartCompleted] = Field(..., title="Parts")
+
+
 class DumpListingsResponse(BaseModel):
     listings: List[Listing] = Field(..., title="Listings")
 

kscale/web/gen/generate.sh

@@ -1,7 +1,10 @@
 #!/bin/sh
 # Install the codegen tool: `pip install datamodel-code-generator`
 
+# for testing against local `www` backend, swap off the production openapi URL
+# openapi_url="http://127.0.0.1:8080/openapi.json"
 openapi_url="https://api.kscale.dev/openapi.json"
+
 curl -s $openapi_url > openapi.json
 datamodel-codegen --input openapi.json --input-file-type openapi --output api.py
 

kscale/web/kernel_images.py

@@ -49,7 +49,7 @@ async def download_kernel_image(artifact_id: str) -> Path:
         headers = {"Authorization": f"Bearer {get_api_key()}", "Accept": "application/octet-stream"}
 
         if not filename.exists():
-            logger.info("Downloading kernel image from %s", artifact_url)
+            logger.info("Downloading kernel image...")
             sha256_hash = hashlib.sha256()
 
             async with httpx.AsyncClient() as client: