From 3281f7abe573a7dc24d4477b620156526387b854 Mon Sep 17 00:00:00 2001 From: Ivan Date: Tue, 19 Nov 2024 09:26:39 -0800 Subject: [PATCH 1/2] Support for uploading kernel images via CLI --- .../listing/ListingArtifactRenderer.tsx | 1 + store/app/crud/artifacts.py | 31 +++++++ store/app/model.py | 82 ++++++++----------- store/settings/environment.py | 2 +- 4 files changed, 67 insertions(+), 49 deletions(-) diff --git a/frontend/src/components/listing/ListingArtifactRenderer.tsx b/frontend/src/components/listing/ListingArtifactRenderer.tsx index af26b793..35873db7 100644 --- a/frontend/src/components/listing/ListingArtifactRenderer.tsx +++ b/frontend/src/components/listing/ListingArtifactRenderer.tsx @@ -20,6 +20,7 @@ const ListingArtifactRenderer = ({ artifact }: Props) => { /> ); case "tgz": + case "kernel": return (
Artifact: + artifact = Artifact.create( + user_id=listing.user_id, + listing_id=listing.id, + name=name, + artifact_type="kernel", + description=description, + ) + + file_data = io.BytesIO(await file.read()) + s3_filename = get_artifact_name(artifact=artifact) + + await asyncio.gather( + self._upload_to_s3( + data=file_data, + name=name, + filename=s3_filename, + content_type="application/octet-stream", + ), + self._add_item(artifact), + ) + return artifact + async def upload_artifact( self, name: str, @@ -271,6 +300,8 @@ async def upload_artifact( match artifact_type: case "image": return await self._upload_image(name, file, listing, description) + case "kernel": + return await self._upload_kernel(name, file, listing, description) case "stl" | "obj" | "ply" | "dae": return await self._upload_mesh(name, file, listing, artifact_type, description) case "urdf" | "mjcf": diff --git a/store/app/model.py b/store/app/model.py index 1d0e6a64..7aeda704 100644 --- a/store/app/model.py +++ b/store/app/model.py @@ -7,6 +7,7 @@ import time from datetime import datetime, timedelta +from pathlib import Path from typing import Literal, Self, cast, get_args from pydantic import BaseModel @@ -173,7 +174,8 @@ def create(cls, user_id: str, source: APIKeySource, permissions: APIKeyPermissio XMLArtifactType = Literal["urdf", "mjcf"] MeshArtifactType = Literal["stl", "obj", "dae", "ply"] CompressedArtifactType = Literal["tgz", "zip"] -ArtifactType = ImageArtifactType | XMLArtifactType | MeshArtifactType | CompressedArtifactType +KernelArtifactType = Literal["kernel"] +ArtifactType = ImageArtifactType | KernelArtifactType | XMLArtifactType | MeshArtifactType | CompressedArtifactType UPLOAD_CONTENT_TYPE_OPTIONS: dict[ArtifactType, set[str]] = { # Image @@ -194,6 +196,13 @@ def create(cls, user_id: str, source: APIKeySource, permissions: APIKeyPermissio "application/x-compressed-tar", }, "zip": {"application/zip"}, + "kernel": { + "application/octet-stream", + "application/x-raw-disk-image", + "application/gzip", + "application/x-gzip", + "binary/octet-stream", + }, } DOWNLOAD_CONTENT_TYPE: dict[ArtifactType, str] = { @@ -210,6 +219,7 @@ def create(cls, user_id: str, source: APIKeySource, permissions: APIKeyPermissio # Compressed "tgz": "application/gzip", "zip": "application/zip", + "kernel": "application/octet-stream", } SizeMapping: dict[ArtifactSize, tuple[int, int]] = { @@ -218,67 +228,43 @@ def create(cls, user_id: str, source: APIKeySource, permissions: APIKeyPermissio } -def get_artifact_type(content_type: str | None, filename: str | None) -> ArtifactType: - """Determines the artifact type from the content type or filename. +def get_artifact_type(content_type: str | None, filename: str) -> ArtifactType: + """Gets the artifact type from the content type and filename.""" + extension = Path(filename).suffix.lower() - Args: - content_type: The content type of the file. - filename: The name of the file. + if extension == ".img": + return "kernel" - Returns: - The artifact type. + if content_type and content_type.startswith("image/"): + return "image" - Raises: - ValueError: If the artifact type cannot be determined. - """ - # Attempts to determine from file extension. - if filename is not None: - extension = filename.split(".")[-1].lower() - if extension in ("png", "jpeg", "jpg", "gif", "webp"): + match extension: + case ".png" | ".jpg" | ".jpeg" | ".gif" | ".webp": return "image" - if extension in ("urdf",): + case ".urdf": return "urdf" - if extension in ("mjcf", "xml"): + case ".mjcf": return "mjcf" - if extension in ("stl",): + case ".stl": return "stl" - if extension in ("obj",): + case ".obj": return "obj" - if extension in ("dae",): + case ".dae": return "dae" - if extension in ("ply",): + case ".ply": return "ply" - if extension in ("tgz", "tar.gz"): + case ".tgz" | ".tar.gz": return "tgz" - if extension in ("zip",): + case ".zip": return "zip" - - # Attempts to determine from content type. - if content_type is not None: - if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["image"]: - return "image" - if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["urdf"]: - return "urdf" - if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["mjcf"]: - return "mjcf" - if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["stl"]: - return "stl" - if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["obj"]: - return "obj" - if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["dae"]: - return "dae" - if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["ply"]: - return "ply" - if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["tgz"]: - return "tgz" - if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["zip"]: - return "zip" - - # Throws a value error if the type cannot be determined. - raise ValueError(f"Unknown content type for file: {filename}") + case _: + raise ValueError(f"Unsupported file extension: {extension}") def get_compression_type(content_type: str | None, filename: str | None) -> CompressedArtifactType: + if filename is None: + raise ValueError("Filename must be provided") + artifact_type = get_artifact_type(content_type, filename) if artifact_type not in (allowed_types := get_args(CompressedArtifactType)): raise ValueError(f"Artifact type {artifact_type} is not compressed; expected one of {allowed_types}") @@ -462,7 +448,7 @@ def get_artifact_name( case "image": height, width = SizeMapping[size] return f"{listing_id}/{artifact_id}/{size}_{height}x{width}_{name}" - case "urdf" | "mjcf" | "stl" | "obj" | "ply" | "dae" | "zip" | "tgz": + case "kernel" | "urdf" | "mjcf" | "stl" | "obj" | "ply" | "dae" | "zip" | "tgz": return f"{listing_id}/{artifact_id}/{name}" case _: raise ValueError(f"Unknown artifact type: {artifact_type}") diff --git a/store/settings/environment.py b/store/settings/environment.py index 17f3fe80..cbdb0b5f 100644 --- a/store/settings/environment.py +++ b/store/settings/environment.py @@ -42,7 +42,7 @@ class ArtifactSettings: large_image_size: tuple[int, int] = field(default=(1536, 1536)) small_image_size: tuple[int, int] = field(default=(256, 256)) min_bytes: int = field(default=16) - max_bytes: int = field(default=1536 * 1536 * 25) + max_bytes: int = field(default=2**30) quality: int = field(default=80) max_concurrent_file_uploads: int = field(default=3) From 3b15970a48289a7b8c44ab671c3982759746028e Mon Sep 17 00:00:00 2001 From: Ivan Date: Tue, 19 Nov 2024 09:47:18 -0800 Subject: [PATCH 2/2] Frontend failure & model changes --- frontend/src/gen/api.ts | 4 +-- store/app/model.py | 60 ++++++++++++++++++++++++++--------------- 2 files changed, 40 insertions(+), 24 deletions(-) diff --git a/frontend/src/gen/api.ts b/frontend/src/gen/api.ts index 35c37369..6df43d17 100644 --- a/frontend/src/gen/api.ts +++ b/frontend/src/gen/api.ts @@ -1925,7 +1925,7 @@ export interface components { /** Name */ name: string; /** Artifact Type */ - artifact_type: "image" | ("urdf" | "mjcf") | ("stl" | "obj" | "dae" | "ply") | ("tgz" | "zip"); + artifact_type: "image" | "kernel" | ("urdf" | "mjcf") | ("stl" | "obj" | "dae" | "ply") | ("tgz" | "zip"); /** Description */ description: string | null; /** Timestamp */ @@ -2459,7 +2459,7 @@ export interface operations { }; header?: never; path: { - artifact_type: "image" | ("urdf" | "mjcf") | ("stl" | "obj" | "dae" | "ply") | ("tgz" | "zip"); + artifact_type: "image" | "kernel" | ("urdf" | "mjcf") | ("stl" | "obj" | "dae" | "ply") | ("tgz" | "zip"); listing_id: string; name: string; }; diff --git a/store/app/model.py b/store/app/model.py index 7aeda704..1e9970fe 100644 --- a/store/app/model.py +++ b/store/app/model.py @@ -7,7 +7,6 @@ import time from datetime import datetime, timedelta -from pathlib import Path from typing import Literal, Self, cast, get_args from pydantic import BaseModel @@ -228,37 +227,54 @@ def create(cls, user_id: str, source: APIKeySource, permissions: APIKeyPermissio } -def get_artifact_type(content_type: str | None, filename: str) -> ArtifactType: - """Gets the artifact type from the content type and filename.""" - extension = Path(filename).suffix.lower() - - if extension == ".img": - return "kernel" - - if content_type and content_type.startswith("image/"): - return "image" +def get_artifact_type(content_type: str | None, filename: str | None) -> ArtifactType: + if filename is not None: + extension = filename.split(".")[-1].lower() + if extension == "img": + return "kernel" + if extension in ("png", "jpeg", "jpg", "gif", "webp"): + return "image" + if extension in ("urdf",): + return "urdf" + if extension in ("mjcf", "xml"): + return "mjcf" + if extension in ("stl",): + return "stl" + if extension in ("obj",): + return "obj" + if extension in ("dae",): + return "dae" + if extension in ("ply",): + return "ply" + if extension in ("tgz", "tar.gz"): + return "tgz" + if extension in ("zip",): + return "zip" - match extension: - case ".png" | ".jpg" | ".jpeg" | ".gif" | ".webp": + # Attempts to determine from content type. + if content_type is not None: + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["kernel"]: + return "kernel" + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["image"]: return "image" - case ".urdf": + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["urdf"]: return "urdf" - case ".mjcf": + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["mjcf"]: return "mjcf" - case ".stl": + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["stl"]: return "stl" - case ".obj": + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["obj"]: return "obj" - case ".dae": + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["dae"]: return "dae" - case ".ply": + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["ply"]: return "ply" - case ".tgz" | ".tar.gz": + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["tgz"]: return "tgz" - case ".zip": + if content_type in UPLOAD_CONTENT_TYPE_OPTIONS["zip"]: return "zip" - case _: - raise ValueError(f"Unsupported file extension: {extension}") + + raise ValueError(f"Unknown content type for file: {filename}") def get_compression_type(content_type: str | None, filename: str | None) -> CompressedArtifactType: