kscale 0.0.13__cp311-cp311-macosx_11_0_arm64.whl

kscale/web/kernels.py

@@ -0,0 +1,207 @@
+"""Utility functions for managing kernel images in K-Scale WWW."""
+
+import hashlib
+import logging
+import shutil
+from pathlib import Path
+
+import aiofiles
+import click
+import httpx
+
+from kscale.utils.checksum import FileChecksum
+from kscale.utils.cli import coro
+from kscale.web.gen.api import SingleArtifactResponse
+from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_key, get_artifact_dir, get_cache_dir
+from kscale.web.www_client import KScaleWWWClient
+
+httpx_logger = logging.getLogger("httpx")
+httpx_logger.setLevel(logging.WARNING)
+
+logger = logging.getLogger(__name__)
+
+ALLOWED_SUFFIXES = {".img"}
+
+
+async def fetch_kernel_image_info(artifact_id: str, cache_dir: Path) -> SingleArtifactResponse:
+    response_path = cache_dir / "response.json"
+    if response_path.exists():
+        return SingleArtifactResponse.model_validate_json(response_path.read_text())
+    async with KScaleWWWClient() as client:
+        response = await client.get_artifact_info(artifact_id)
+    response_path.write_text(response.model_dump_json())
+    return response
+
+
+async def download_kernel_image(artifact_id: str) -> Path:
+    """Download a kernel image artifact."""
+    cache_dir = get_artifact_dir(artifact_id)
+    try:
+        kernel_image_info = await fetch_kernel_image_info(artifact_id, cache_dir)
+        artifact_url = kernel_image_info.urls.large
+
+        original_name = Path(artifact_url).name
+        if not original_name.endswith(".img"):
+            filename = cache_dir / f"{original_name}.img"
+        else:
+            filename = cache_dir / original_name
+
+        headers = {
+            "Authorization": f"Bearer {get_api_key()}",
+            "Accept": "application/octet-stream",
+        }
+
+        if not filename.exists():
+            logger.info("Downloading kernel image...")
+            sha256_hash = hashlib.sha256()
+
+            async with httpx.AsyncClient() as client:
+                async with client.stream("GET", artifact_url, headers=headers) as response:
+                    response.raise_for_status()
+
+                    async with aiofiles.open(filename, "wb") as f:
+                        async for chunk in response.aiter_bytes():
+                            FileChecksum.update_hash(sha256_hash, chunk)
+                            await f.write(chunk)
+
+            logger.info("Kernel image downloaded to %s", filename)
+        else:
+            logger.info("Kernel image already cached at %s", filename)
+
+        return filename
+
+    except httpx.RequestError:
+        logger.exception("Failed to fetch kernel image")
+        raise
+
+
+async def show_kernel_image_info(artifact_id: str) -> None:
+    """Show information about a kernel image artifact."""
+    try:
+        kernel_image_info = await fetch_kernel_image_info(artifact_id, get_artifact_dir(artifact_id))
+        logger.info("Kernel Image Artifact ID: %s", kernel_image_info.artifact_id)
+        logger.info("Kernel Image URL: %s", kernel_image_info.urls.large)
+    except httpx.RequestError:
+        logger.exception("Failed to fetch kernel image info")
+        raise
+
+
+async def remove_local_kernel_image(artifact_id: str) -> None:
+    """Remove local cache of a kernel image artifact."""
+    try:
+        if artifact_id.lower() == "all":
+            cache_dir = get_cache_dir()
+            if cache_dir.exists():
+                logger.info("Removing all local caches at %s", cache_dir)
+                shutil.rmtree(cache_dir)
+            else:
+                logger.error("No local caches found")
+        else:
+            artifact_dir = get_artifact_dir(artifact_id)
+            if artifact_dir.exists():
+                logger.info("Removing local cache at %s", artifact_dir)
+                shutil.rmtree(artifact_dir)
+            else:
+                logger.error("No local cache found for artifact %s", artifact_id)
+
+    except Exception:
+        logger.error("Failed to remove local cache")
+        raise
+
+
+async def upload_kernel_image(
+    listing_id: str,
+    image_path: Path,
+    upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT,
+) -> SingleArtifactResponse:
+    """Upload a kernel image."""
+    if image_path.suffix.lower() not in ALLOWED_SUFFIXES:
+        raise ValueError("Invalid file type. Must be one of: %s", ALLOWED_SUFFIXES)
+
+    if not image_path.exists():
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+
+    checksum, file_size = await FileChecksum.calculate(str(image_path))
+    logger.info("Uploading kernel image: %s", image_path)
+    logger.info("File name: %s", image_path.name)
+    logger.info("File size: %.1f MB", file_size / 1024 / 1024)
+
+    async with KScaleWWWClient(upload_timeout=upload_timeout) as client:
+        presigned_data = await client.get_presigned_url(
+            listing_id=listing_id,
+            file_name=image_path.name,
+            checksum=checksum,
+        )
+
+        logger.info("Starting upload...")
+        async with httpx.AsyncClient() as http_client:
+            logger.info("Reading file content into memory...")
+            async with aiofiles.open(image_path, "rb") as f:
+                contents = await f.read()
+
+            logger.info("Uploading file content to %s", presigned_data["upload_url"])
+            response = await http_client.put(
+                presigned_data["upload_url"],
+                content=contents,
+                headers={"Content-Type": "application/x-raw-disk-image"},
+                timeout=upload_timeout,
+            )
+            response.raise_for_status()
+
+        artifact_response: SingleArtifactResponse = await client.get_artifact_info(presigned_data["artifact_id"])
+        logger.info("Uploaded artifact: %s", artifact_response.artifact_id)
+        return artifact_response
+
+
+async def upload_kernel_image_cli(
+    listing_id: str, image_path: Path, upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT
+) -> SingleArtifactResponse:
+    """CLI wrapper for upload_kernel_image."""
+    response = await upload_kernel_image(listing_id, image_path, upload_timeout=upload_timeout)
+    return response
+
+
+@click.group()
+def cli() -> None:
+    """K-Scale Kernel Image CLI tool."""
+    pass
+
+
+@cli.command()
+@click.argument("artifact_id")
+@coro
+async def download(artifact_id: str) -> None:
+    """Download a kernel image artifact."""
+    await download_kernel_image(artifact_id)
+
+
+@cli.command()
+@click.argument("artifact_id")
+@coro
+async def info(artifact_id: str) -> None:
+    """Show information about a kernel image artifact."""
+    await show_kernel_image_info(artifact_id)
+
+
+@cli.command("remove-local")
+@click.argument("artifact_id")
+@coro
+async def remove_local(artifact_id: str) -> None:
+    """Remove local cache of a kernel image artifact."""
+    await remove_local_kernel_image(artifact_id)
+
+
+@cli.command()
+@click.argument("listing_id")
+@click.argument("image_path", type=click.Path(exists=True, path_type=Path))
+@click.option(
+    "--upload-timeout", type=float, default=DEFAULT_UPLOAD_TIMEOUT, help="Timeout in seconds for upload operations"
+)
+@coro
+async def upload(listing_id: str, image_path: Path, upload_timeout: float) -> None:
+    """Upload a kernel image artifact."""
+    await upload_kernel_image_cli(listing_id, image_path, upload_timeout=upload_timeout)
+
+
+if __name__ == "__main__":
+    cli()

kscale/web/krec.py

@@ -0,0 +1,175 @@
+"""Utility functions for managing K-Recs in K-Scale WWW."""
+
+import asyncio
+import json
+import logging
+from pathlib import Path
+
+import aiofiles
+import click
+import httpx
+import krec
+
+from kscale.utils.cli import coro
+from kscale.web.gen.api import UploadKRecRequest
+from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_key, get_artifact_dir
+from kscale.web.www_client import KScaleWWWClient
+
+logger = logging.getLogger(__name__)
+
+
+async def upload_krec(
+    robot_id: str,
+    file_path: Path,
+    description: str | None = None,
+    upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT,
+) -> str:
+    if not file_path.exists():
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    file_size = file_path.stat().st_size
+    logger.info("Uploading K-Rec: %s", file_path)
+    logger.info("File name: %s", file_path.name)
+    logger.info("File size: %.1f MB", file_size / 1024 / 1024)
+
+    if not file_path.suffix.lower() == ".krec":
+        logger.warning("File extension is not .krec - are you sure this is a valid K-Rec file?")
+
+    try:
+        krec.KRec.load(str(file_path.resolve()))
+    except Exception as e:
+        raise ValueError(f"Failed to load K-Rec from {file_path} - are you sure this is a valid K-Rec file?") from e
+
+    async with KScaleWWWClient(upload_timeout=upload_timeout) as client:
+        create_response = await client.create_krec(
+            UploadKRecRequest(
+                robot_id=robot_id,
+                name=file_path.name,
+                description=description,
+            )
+        )
+
+        logger.info("Initialized K-Rec upload with ID: %s", create_response["krec_id"])
+        logger.info("Starting upload...")
+        async with httpx.AsyncClient() as http_client:
+            logger.info("Reading file content into memory...")
+            async with aiofiles.open(file_path, "rb") as f:
+                contents = await f.read()
+
+            logger.info("Uploading file content to %s", create_response["upload_url"])
+            response = await http_client.put(
+                create_response["upload_url"],
+                content=contents,
+                headers={"Content-Type": "video/x-matroska"},
+                timeout=upload_timeout,
+            )
+            response.raise_for_status()
+
+        logger.info("Successfully uploaded K-Rec: %s", create_response["krec_id"])
+        return create_response["krec_id"]
+
+
+def upload_krec_sync(robot_id: str, file_path: Path, description: str | None = None) -> str:
+    return asyncio.run(upload_krec(robot_id, file_path, description))
+
+
+async def fetch_krec_info(krec_id: str, cache_dir: Path) -> dict:
+    """Fetch K-Rec info from the server or cache."""
+    response_path = cache_dir / "response.json"
+    if response_path.exists():
+        return json.loads(response_path.read_text())
+
+    async with KScaleWWWClient() as client:
+        try:
+            response = await client.get_krec_info(krec_id)
+
+            if not response:
+                raise ValueError(f"Empty response from server for K-Rec ID: {krec_id}")
+
+            response_path.write_text(json.dumps(response))
+            return response
+        except Exception as e:
+            logger.error("Error fetching K-Rec info: %s", str(e))
+            raise
+
+
+async def download_krec(krec_id: str) -> Path:
+    """Download a K-Rec file."""
+    cache_dir = get_artifact_dir(krec_id)
+
+    try:
+        krec_info = await fetch_krec_info(krec_id, cache_dir)
+
+        if not isinstance(krec_info, dict):
+            logger.error("Unexpected response type: %s", type(krec_info))
+            raise ValueError(f"Invalid response format for K-Rec ID: {krec_id}")
+
+        if "url" not in krec_info or "filename" not in krec_info:
+            logger.error("Response missing required fields: %s", krec_info)
+            raise ValueError(f"Invalid response format for K-Rec ID: {krec_id}")
+
+        download_url = krec_info["url"]
+        filename = krec_info["filename"]
+        full_path = cache_dir / filename
+
+        if full_path.exists():
+            logger.info("K-Rec already cached at %s", full_path)
+            return full_path
+
+        logger.info("Downloading K-Rec %s to %s", krec_id, full_path)
+
+        api_key = get_api_key()
+        headers = {"Authorization": f"Bearer {api_key}", "Accept": "application/octet-stream"}
+
+        async with httpx.AsyncClient() as client:
+            async with client.stream("GET", download_url, headers=headers) as response:
+                response.raise_for_status()
+
+                with open(full_path, "wb") as f:
+                    async for chunk in response.aiter_bytes():
+                        f.write(chunk)
+
+        logger.info("Successfully downloaded K-Rec to %s", full_path)
+        return full_path
+
+    except httpx.RequestError as e:
+        logger.exception("Failed to fetch K-Rec: %s", str(e))
+        raise
+    except Exception as e:
+        logger.exception("Unexpected error: %s", str(e))
+        raise
+
+
+def download_krec_sync(krec_id: str) -> Path:
+    """Sync wrapper for download_krec."""
+    return asyncio.run(download_krec(krec_id))
+
+
+@click.group(name="krec")
+def cli() -> None:
+    """K-Scale K-Rec management commands."""
+    pass
+
+
+@cli.command()
+@click.argument("robot_id")
+@click.argument("file_path", type=click.Path(exists=True, path_type=Path))
+@click.option("--description", "-d", help="Description of the K-Rec")
+@coro
+async def upload(robot_id: str, file_path: Path, description: str | None = None) -> None:
+    """Upload a K-Rec file."""
+    krec_id = await upload_krec(robot_id, file_path, description)
+    click.echo(f"Successfully uploaded K-Rec: {krec_id}")
+
+
+@cli.command()
+@click.argument("krec_id")
+@coro
+async def download(krec_id: str) -> None:
+    """Download a K-Rec file."""
+    file_path = await download_krec(krec_id)
+    click.echo(f"Successfully downloaded K-Rec to: {file_path}")
+
+
+if __name__ == "__main__":
+    cli()

kscale/web/pybullet.py

@@ -0,0 +1,188 @@
+"""Simple script to interact with a URDF in PyBullet."""
+
+import itertools
+import logging
+import math
+import time
+
+import click
+
+from kscale.artifacts import PLANE_URDF_PATH
+from kscale.utils.cli import coro
+from kscale.web.urdf import download_urdf
+
+logger = logging.getLogger(__name__)
+
+
+@click.command(help="Show a URDF in PyBullet")
+@click.argument("listing_id")
+@click.option("--dt", type=float, default=0.01, help="Time step")
+@click.option("-n", "--hide-gui", is_flag=True, help="Hide the GUI")
+@click.option("--no-merge", is_flag=True, help="Do not merge fixed links")
+@click.option("--hide-origin", is_flag=True, help="Do not show the origin")
+@click.option("--show-inertia", is_flag=True, help="Visualizes the inertia frames")
+@click.option("--see-thru", is_flag=True, help="Use see-through mode")
+@click.option("--show-collision", is_flag=True, help="Show collision meshes")
+@coro
+async def cli(
+    listing_id: str,
+    dt: float,
+    hide_gui: bool,
+    no_merge: bool,
+    hide_origin: bool,
+    show_inertia: bool,
+    see_thru: bool,
+    show_collision: bool,
+) -> None:
+    # Gets the URDF path.
+    urdf_dir = await download_urdf(listing_id)
+    urdf_path = next(urdf_dir.glob("*.urdf"), None)
+    if urdf_path is None:
+        raise ValueError(f"No URDF found in {urdf_dir}")
+
+    try:
+        import pybullet as p  # type: ignore[import-not-found]
+    except ImportError:
+        raise ImportError("pybullet is required to run this script")
+
+    # Connect to PyBullet.
+    p.connect(p.GUI)
+    p.setGravity(0, 0, -9.81)
+    p.setRealTimeSimulation(0)
+
+    # Turn off panels.
+    if hide_gui:
+        p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
+    p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0)
+    p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0)
+    p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0)
+
+    # Enable mouse picking.
+    p.configureDebugVisualizer(p.COV_ENABLE_MOUSE_PICKING, 1)
+
+    # Loads the floor plane.
+    floor = p.loadURDF(str(PLANE_URDF_PATH))
+
+    # Load the robot URDF.
+    start_position = [0.0, 0.0, 1.0]
+    start_orientation = p.getQuaternionFromEuler([0.0, 0.0, 0.0])
+    flags = p.URDF_USE_INERTIA_FROM_FILE
+    if not no_merge:
+        flags |= p.URDF_MERGE_FIXED_LINKS
+    robot = p.loadURDF(str(urdf_path), start_position, start_orientation, flags=flags, useFixedBase=0)
+
+    # Display collision meshes as separate object.
+    if show_collision:
+        collision_flags = p.URDF_USE_INERTIA_FROM_FILE | p.URDF_USE_SELF_COLLISION_EXCLUDE_ALL_PARENTS
+        collision = p.loadURDF(str(urdf_path), start_position, start_orientation, flags=collision_flags, useFixedBase=0)
+
+        # Make collision shapes semi-transparent.
+        joint_ids = [i for i in range(p.getNumJoints(collision))] + [-1]
+        for i in joint_ids:
+            p.changeVisualShape(collision, i, rgbaColor=[1, 0, 0, 0.5])
+
+    # Initializes physics parameters.
+    p.changeDynamics(floor, -1, lateralFriction=1, spinningFriction=-1, rollingFriction=-1)
+    p.setPhysicsEngineParameter(fixedTimeStep=dt, maxNumCmdPer1ms=1000)
+
+    # Shows the origin of the robot.
+    if not hide_origin:
+        p.addUserDebugLine([0, 0, 0], [0.1, 0, 0], [1, 0, 0], parentObjectUniqueId=robot, parentLinkIndex=-1)
+        p.addUserDebugLine([0, 0, 0], [0, 0.1, 0], [0, 1, 0], parentObjectUniqueId=robot, parentLinkIndex=-1)
+        p.addUserDebugLine([0, 0, 0], [0, 0, 0.1], [0, 0, 1], parentObjectUniqueId=robot, parentLinkIndex=-1)
+
+    # Make the robot see-through.
+    joint_ids = [i for i in range(p.getNumJoints(robot))] + [-1]
+    if see_thru:
+        for i in joint_ids:
+            p.changeVisualShape(robot, i, rgbaColor=[1, 1, 1, 0.5])
+
+    def draw_box(pt: list[list[float]], color: tuple[float, float, float], obj_id: int, link_id: int) -> None:
+        assert len(pt) == 8
+        assert all(len(p) == 3 for p in pt)
+
+        mapping = [1, 3, 0, 2]
+        for i in range(4):
+            p.addUserDebugLine(pt[i], pt[i + 4], color, 1, parentObjectUniqueId=obj_id, parentLinkIndex=link_id)
+            p.addUserDebugLine(pt[i], pt[mapping[i]], color, 1, parentObjectUniqueId=obj_id, parentLinkIndex=link_id)
+            p.addUserDebugLine(
+                pt[i + 4], pt[mapping[i] + 4], color, 1, parentObjectUniqueId=obj_id, parentLinkIndex=link_id
+            )
+
+    # Shows bounding boxes around each part of the robot representing the inertia frame.
+    if show_inertia:
+        for i in joint_ids:
+            dynamics_info = p.getDynamicsInfo(robot, i)
+            mass = dynamics_info[0]
+            if mass <= 0:
+                continue
+            inertia = dynamics_info[2]
+            ixx = inertia[0]
+            iyy = inertia[1]
+            izz = inertia[2]
+            box_scale_x = 0.5 * math.sqrt(6 * (izz + iyy - ixx) / mass)
+            box_scale_y = 0.5 * math.sqrt(6 * (izz + ixx - iyy) / mass)
+            box_scale_z = 0.5 * math.sqrt(6 * (ixx + iyy - izz) / mass)
+
+            half_extents = [box_scale_x, box_scale_y, box_scale_z]
+            pt = [
+                [x, y, z]
+                for x, y, z in itertools.product(
+                    [-half_extents[0], half_extents[0]],
+                    [-half_extents[1], half_extents[1]],
+                    [-half_extents[2], half_extents[2]],
+                )
+            ]
+            draw_box(pt, (1, 0, 0), robot, i)
+
+    # Show joint controller.
+    joints: dict[str, int] = {}
+    controls: dict[str, float] = {}
+    for i in range(p.getNumJoints(robot)):
+        joint_info = p.getJointInfo(robot, i)
+        name = joint_info[1].decode("utf-8")
+        joint_type = joint_info[2]
+        joints[name] = i
+        if joint_type == p.JOINT_PRISMATIC:
+            joint_min, joint_max = joint_info[8:10]
+            controls[name] = p.addUserDebugParameter(name, joint_min, joint_max, 0.0)
+        elif joint_type == p.JOINT_REVOLUTE:
+            joint_min, joint_max = joint_info[8:10]
+            controls[name] = p.addUserDebugParameter(name, joint_min, joint_max, 0.0)
+
+    # Run the simulation until the user closes the window.
+    last_time = time.time()
+    prev_control_values = {k: 0.0 for k in controls}
+    while p.isConnected():
+        # Reset the simulation if "r" was pressed.
+        keys = p.getKeyboardEvents()
+        if ord("r") in keys and keys[ord("r")] & p.KEY_WAS_TRIGGERED:
+            p.resetBasePositionAndOrientation(robot, start_position, start_orientation)
+            p.setJointMotorControlArray(
+                robot,
+                range(p.getNumJoints(robot)),
+                p.POSITION_CONTROL,
+                targetPositions=[0] * p.getNumJoints(robot),
+            )
+
+        # Set joint positions.
+        for k, v in controls.items():
+            try:
+                target_position = p.readUserDebugParameter(v)
+                if target_position != prev_control_values[k]:
+                    prev_control_values[k] = target_position
+                    p.setJointMotorControl2(robot, joints[k], p.POSITION_CONTROL, target_position)
+            except p.error:
+                logger.debug("Failed to set joint %s", k)
+                pass
+
+        # Step simulation.
+        p.stepSimulation()
+        cur_time = time.time()
+        time.sleep(max(0, dt - (cur_time - last_time)))
+        last_time = cur_time
+
+
+if __name__ == "__main__":
+    # python -m kscale.web.pybullet
+    cli()