kscale 0.0.13__cp311-cp311-macosx_11_0_arm64.whl

kscale/web/urdf.py

@@ -0,0 +1,185 @@
+"""Utility functions for managing artifacts in K-Scale WWW."""
+
+import logging
+import shutil
+import tarfile
+from pathlib import Path
+
+import click
+import httpx
+import requests
+
+from kscale.utils.cli import coro
+from kscale.web.gen.api import SingleArtifactResponse, UploadArtifactResponse
+from kscale.web.utils import get_api_key, get_artifact_dir, get_cache_dir
+from kscale.web.www_client import KScaleWWWClient
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+ALLOWED_SUFFIXES = {
+    ".urdf",
+    ".mjcf",
+    ".stl",
+    ".obj",
+    ".dae",
+    ".png",
+    ".jpg",
+    ".jpeg",
+}
+
+
+async def fetch_urdf_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_artifact(artifact_url: str, cache_dir: Path) -> Path:
+    filename = cache_dir / Path(artifact_url).name
+    headers = {
+        "Authorization": f"Bearer {get_api_key()}",
+    }
+
+    if not filename.exists():
+        logger.info("Downloading artifact from %s", artifact_url)
+
+        async with httpx.AsyncClient() as client:
+            response = await client.get(artifact_url, headers=headers)
+            response.raise_for_status()
+            filename.write_bytes(response.content)
+            logger.info("Artifact downloaded to %s", filename)
+    else:
+        logger.info("Artifact already cached at %s", filename)
+
+    # Extract the .tgz file
+    extract_dir = cache_dir / filename.stem
+    if not extract_dir.exists():
+        logger.info("Extracting %s to %s", filename, extract_dir)
+        with tarfile.open(filename, "r:gz") as tar:
+            tar.extractall(path=extract_dir)
+    else:
+        logger.info("Artifact already extracted at %s", extract_dir)
+
+    return extract_dir
+
+
+def create_tarball(folder_path: Path, output_filename: str, cache_dir: Path) -> Path:
+    tarball_path = cache_dir / output_filename
+    with tarfile.open(tarball_path, "w:gz") as tar:
+        for file_path in folder_path.rglob("*"):
+            if file_path.is_file() and file_path.suffix.lower() in ALLOWED_SUFFIXES:
+                tar.add(file_path, arcname=file_path.relative_to(folder_path))
+                logger.info("Added %s to tarball", file_path)
+            else:
+                logger.warning("Skipping %s", file_path)
+    logger.info("Created tarball %s", tarball_path)
+    return tarball_path
+
+
+async def download_urdf(artifact_id: str) -> Path:
+    cache_dir = get_artifact_dir(artifact_id)
+    try:
+        urdf_info = await fetch_urdf_info(artifact_id, cache_dir)
+        artifact_url = urdf_info.urls.large
+        return await download_artifact(artifact_url, cache_dir)
+
+    except requests.RequestException:
+        logger.exception("Failed to fetch URDF info")
+        raise
+
+
+async def show_urdf_info(artifact_id: str) -> None:
+    try:
+        urdf_info = await fetch_urdf_info(artifact_id, get_artifact_dir(artifact_id))
+        logger.info("URDF Artifact ID: %s", urdf_info.artifact_id)
+        logger.info("URDF URL: %s", urdf_info.urls.large)
+    except requests.RequestException:
+        logger.exception("Failed to fetch URDF info")
+        raise
+
+
+async def remove_local_urdf(artifact_id: str) -> None:
+    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_urdf(listing_id: str, root_dir: Path) -> UploadArtifactResponse:
+    tarball_path = create_tarball(root_dir, "robot.tgz", get_artifact_dir(listing_id))
+
+    async with KScaleWWWClient() as client:
+        response = await client.upload_artifact(listing_id, str(tarball_path))
+
+    logger.info("Uploaded artifacts: %s", [artifact.artifact_id for artifact in response.artifacts])
+    return response
+
+
+async def upload_urdf_cli(listing_id: str, root_dir: Path) -> UploadArtifactResponse:
+    response = await upload_urdf(listing_id, root_dir)
+    return response
+
+
+@click.group()
+def cli() -> None:
+    """K-Scale URDF Store CLI tool."""
+    pass
+
+
+@cli.command()
+@click.argument("artifact_id")
+@coro
+async def download(artifact_id: str) -> None:
+    """Download a URDF artifact."""
+    await download_urdf(artifact_id)
+
+
+@cli.command()
+@click.argument("artifact_id")
+@coro
+async def info(artifact_id: str) -> None:
+    """Show information about a URDF artifact."""
+    await show_urdf_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 URDF artifact."""
+    await remove_local_urdf(artifact_id)
+
+
+@cli.command()
+@click.argument("listing_id")
+@click.argument("root_dir", type=click.Path(exists=True, path_type=Path))
+@coro
+async def upload(listing_id: str, root_dir: Path) -> None:
+    """Upload a URDF artifact."""
+    await upload_urdf_cli(listing_id, root_dir)
+
+
+if __name__ == "__main__":
+    # python -m kscale.web.urdf
+    cli()

kscale/web/utils.py

@@ -0,0 +1,48 @@
+"""Utility functions for interacting with the K-Scale WWW API."""
+
+import os
+from pathlib import Path
+
+from kscale.conf import Settings
+
+DEFAULT_UPLOAD_TIMEOUT = 300.0  # 5 minutes
+
+
+def get_api_root() -> str:
+    """Returns the base URL for the K-Scale WWW API.
+
+    This can be overridden when targetting a different server.
+
+    Returns:
+        The base URL for the K-Scale WWW API.
+    """
+    return os.getenv("KSCALE_API_ROOT", "https://api.kscale.dev")
+
+
+def get_api_key() -> str:
+    """Returns the API key for the K-Scale WWW API.
+
+    Returns:
+        The API key for the K-Scale WWW API.
+    """
+    api_key = Settings.load().www.api_key
+    if api_key is None:
+        api_key = os.getenv("KSCALE_API_KEY")
+    if not api_key:
+        raise ValueError(
+            "API key not found! Get one here and set it as the `KSCALE_API_KEY` environment variable or in your "
+            "config file: https://kscale.dev/keys"
+        )
+    return api_key
+
+
+def get_cache_dir() -> Path:
+    """Returns the cache directory for artifacts."""
+    return Path(Settings.load().www.cache_dir).expanduser().resolve()
+
+
+def get_artifact_dir(artifact_id: str) -> Path:
+    """Returns the directory for a specific artifact."""
+    cache_dir = get_cache_dir() / artifact_id
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir

kscale/web/www_client.py

@@ -0,0 +1,134 @@
+"""Defines a typed client for the K-Scale WWW API."""
+
+import logging
+from pathlib import Path
+from types import TracebackType
+from typing import Any, Dict, Type
+from urllib.parse import urljoin
+
+import httpx
+from pydantic import BaseModel
+
+from kscale.web.gen.api import (
+    BodyAddListingListingsAddPost,
+    NewListingResponse,
+    SingleArtifactResponse,
+    UploadArtifactResponse,
+    UploadKRecRequest,
+)
+from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_key, get_api_root
+
+logger = logging.getLogger(__name__)
+
+
+class KScaleWWWClient:
+    def __init__(self, base_url: str = get_api_root(), upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT) -> None:
+        self.base_url = base_url
+        self.upload_timeout = upload_timeout
+        self._client: httpx.AsyncClient | None = None
+
+    @property
+    def client(self) -> httpx.AsyncClient:
+        if self._client is None:
+            self._client = httpx.AsyncClient(
+                base_url=self.base_url,
+                headers={"Authorization": f"Bearer {get_api_key()}"},
+                timeout=httpx.Timeout(30.0),
+            )
+        return self._client
+
+    async def _request(
+        self,
+        method: str,
+        endpoint: str,
+        *,
+        params: Dict[str, Any] | None = None,
+        data: BaseModel | None = None,
+        files: Dict[str, Any] | None = None,
+    ) -> Dict[str, Any]:
+        url = urljoin(self.base_url, endpoint)
+        kwargs: Dict[str, Any] = {"params": params}
+
+        if data:
+            kwargs["json"] = data.dict(exclude_unset=True)
+        if files:
+            kwargs["files"] = files
+
+        response = await self.client.request(method, url, **kwargs)
+
+        if response.is_error:
+            logger.error("Error response from K-Scale: %s", response.text)
+        response.raise_for_status()
+        return response.json()
+
+    async def get_artifact_info(self, artifact_id: str) -> SingleArtifactResponse:
+        data = await self._request("GET", f"/artifacts/info/{artifact_id}")
+        return SingleArtifactResponse(**data)
+
+    async def upload_artifact(self, listing_id: str, file_path: str) -> UploadArtifactResponse:
+        file_name = Path(file_path).name
+        with open(file_path, "rb") as f:
+            files = {"files": (file_name, f, "application/gzip")}
+            data = await self._request("POST", f"/artifacts/upload/{listing_id}", files=files)
+        return UploadArtifactResponse(**data)
+
+    async def create_listing(self, request: BodyAddListingListingsAddPost) -> NewListingResponse:
+        data = await self._request("POST", "/listings", data=request)
+        return NewListingResponse(**data)
+
+    async def create_krec(self, request: UploadKRecRequest) -> dict:
+        """Create a new K-Rec upload and get the presigned URL."""
+        return await self._request(
+            "POST",
+            "/krecs/upload",
+            data=request,
+        )
+
+    async def close(self) -> None:
+        if self._client is not None:
+            await self._client.aclose()
+            self._client = None
+
+    async def __aenter__(self) -> "KScaleWWWClient":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.close()
+
+    async def upload_to_presigned_url(self, url: str, file_path: str) -> None:
+        """Upload a file using a presigned URL."""
+        with open(file_path, "rb") as f:
+            async with httpx.AsyncClient(timeout=httpx.Timeout(timeout=self.upload_timeout)) as client:
+                response = await client.put(url, content=f.read(), headers={"Content-Type": "application/octet-stream"})
+                response.raise_for_status()
+
+    async def get_presigned_url(self, listing_id: str, file_name: str, checksum: str | None = None) -> dict:
+        """Get a presigned URL for uploading an artifact."""
+        params = {"filename": file_name}
+        if checksum:
+            params["checksum"] = checksum
+        return await self._request("POST", f"/artifacts/presigned/{listing_id}", params=params)
+
+    async def get_krec_info(self, krec_id: str) -> dict:
+        """Get information about a K-Rec."""
+        logger.info("Getting K-Rec info for ID: %s", krec_id)
+        try:
+            data = await self._request("GET", f"/krecs/download/{krec_id}")
+            if not isinstance(data, dict):
+                logger.error("Server returned unexpected type: %s", type(data))
+                logger.error("Response data: %s", data)
+                raise ValueError(f"Server returned {type(data)} instead of dictionary")
+
+            return {
+                "url": data.get("url"),
+                "filename": data.get("filename"),
+                "checksum": data.get("checksum"),
+            }
+        except Exception as e:
+            logger.error("Failed to get K-Rec info: %s", str(e))
+            raise

kscale-0.0.13.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Benjamin Bolte
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

kscale-0.0.13.dist-info/METADATA

@@ -0,0 +1,55 @@
+Metadata-Version: 2.1
+Name: kscale
+Version: 0.0.13
+Summary: The kscale project
+Home-page: https://github.com/kscalelabs/kscale
+Author: Benjamin Bolte
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: omegaconf
+Requires-Dist: email_validator
+Requires-Dist: httpx
+Requires-Dist: requests
+Requires-Dist: pydantic
+Requires-Dist: click
+Requires-Dist: aiofiles
+Requires-Dist: krec
+Provides-Extra: dev
+Requires-Dist: black; extra == "dev"
+Requires-Dist: darglint; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: datamodel-code-generator; extra == "dev"
+
+<p align="center">
+  <picture>
+    <img alt="K-Scale Open Source Robotics" src="https://media.kscale.dev/kscale-open-source-header.png" style="max-width: 100%;">
+  </picture>
+</p>
+
+<div align="center">
+
+[![License](https://img.shields.io/badge/license-MIT-green)](https://github.com/kscalelabs/ksim/blob/main/LICENSE)
+[![Discord](https://img.shields.io/discord/1224056091017478166)](https://discord.gg/k5mSvCkYQh)
+[![Wiki](https://img.shields.io/badge/wiki-humanoids-black)](https://humanoids.wiki)
+<br />
+[![python](https://img.shields.io/badge/-Python_3.11-blue?logo=python&logoColor=white)](https://github.com/pre-commit/pre-commit)
+[![black](https://img.shields.io/badge/Code%20Style-Black-black.svg?labelColor=gray)](https://black.readthedocs.io/en/stable/)
+[![ruff](https://img.shields.io/badge/Linter-Ruff-red.svg?labelColor=gray)](https://github.com/charliermarsh/ruff)
+<br />
+[![Python Checks](https://github.com/kscalelabs/kscale/actions/workflows/test.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/test.yml)
+[![Publish Python Package](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml)
+
+</div>
+
+# K-Scale Command Line Interface
+
+This is a command line tool for interacting with various services provided by K-Scale Labs. For more information, see the [documentation](https://docs.kscale.dev/pkg/intro).
+
+## Installation
+
+```bash
+pip install kscale
+```

kscale-0.0.13.dist-info/RECORD

@@ -0,0 +1,31 @@
+kscale/requirements.txt,sha256=zj7Rg5vE2Fl3ao66K9jmOlhv5_8SybHvHK5-4xOGZBs,141
+kscale/conf.py,sha256=iI6eWgzIhzmqqy9EhVh6MbsmNpqCbTlSVdLEYfxyjvI,1416
+kscale/__init__.py,sha256=FOc1X2Sjq58lZ0IMP9e0P4HQ2MM_7yBHk-KbuvApm-0,173
+kscale/api.py,sha256=314tgy4e9fbQWsKPYLyofZyiKfKr9RmvstldXs2jDvY,322
+kscale/cli.py,sha256=kfu-5dCr0DIEDEnF80tg5md56bEI2GsyA7mEuad1cDw,655
+kscale/requirements-dev.txt,sha256=WI7-ea4IRJakmqVMN8QKhOsDGrghwtvk03aIsFaNSIw,130
+kscale/rust.cpython-311-darwin.so,sha256=_mrfXiTsFL75KVnkmr0ikDDVLuftiP7lnleR1iQkVrE,431968
+kscale/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kscale/artifacts/plane.obj,sha256=x59-IIrWpLjhotChiqT2Ul6U8s0RcHkaEeUZb4KXL1c,348
+kscale/artifacts/__init__.py,sha256=RK8wdybtCJPgdLLJ8R8-YMi1Ph5ojqAKVJZowHONtgo,232
+kscale/artifacts/plane.urdf,sha256=LCiTk14AyTHjkZ1jvsb0hNaEaJUxDb8Z1JjsgpXu3YM,819
+kscale/web/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kscale/web/pybullet.py,sha256=_at_1yf63Fzy7LmXbR-4fvdM9LqPmzYiMniqhKU-BtE,7395
+kscale/web/api.py,sha256=f7_BPwMnByXSBDE1Rzuhr7IIqoIBY2wqDGoqFLQ_0ng,3809
+kscale/web/www_client.py,sha256=Vu_rf2NuJ7fSiqTEb1B71OoTurdNrCtSVKc7K8AIk24,4963
+kscale/web/utils.py,sha256=mXfrtH2TuxXkSF01QpcUyC7OHAszQPsiFYGZcHHgTDk,1337
+kscale/web/urdf.py,sha256=rJjQBG3q5HBRxebNAccOqD-rXMS-2fpNdh8k5-6p4p0,5821
+kscale/web/kernels.py,sha256=ddp8IEFSLS9raETQZIphDjxbrj9Jd0gfeQ6Fad0ee-Q,7242
+kscale/web/krec.py,sha256=k9AhBHwLoxRDRUSllgQ88C_xMzwHCY4mHRrBsQFP1EA,5910
+kscale/web/gen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kscale/web/gen/api.py,sha256=JqNit9Q6zUuWz-SrWjcKx3BQIzSmL50r0BT1CfT2hbo,23092
+kscale/utils/checksum.py,sha256=jt6QmmQND9zrOEnUtOfZpLYROhgto4Gh3OpdUWk4tZA,1093
+kscale/utils/api_base.py,sha256=Kk_WtRDdJHmOg6NtHmVxVrcfARSUkhfr29ypLch_pO0,112
+kscale/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kscale/utils/cli.py,sha256=JoaY5x5SdUx97KmMM9j5AjRRUqqrTlJ9qVckZptEsYA,827
+kscale-0.0.13.dist-info/RECORD,,
+kscale-0.0.13.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+kscale-0.0.13.dist-info/WHEEL,sha256=aDf-IWQmq1RVXR_5Ck1qPm3e7xxy4ZeSSdW4ZZA07R8,109
+kscale-0.0.13.dist-info/entry_points.txt,sha256=c2pxylFuUJe3aQ_GLEgZrrBW98PPNkiH8Lg0NL_vnrA,42
+kscale-0.0.13.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
+kscale-0.0.13.dist-info/METADATA,sha256=-Nmkr2zrj4ydnMtEW2miV_dqgoyY5PJuiirwUrIRsF8,2174

kscale-0.0.13.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.6.0)
+Root-Is-Purelib: false
+Tag: cp311-cp311-macosx_11_0_arm64
+

kscale-0.0.13.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+kscale = kscale.cli:cli

kscale-0.0.13.dist-info/top_level.txt

@@ -0,0 +1 @@
+kscale