kscale 0.0.13__cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl

Sign up to get free protection for your applications and to get access to all the features.
kscale/web/urdf.py ADDED
@@ -0,0 +1,185 @@
1
+ """Utility functions for managing artifacts in K-Scale WWW."""
2
+
3
+ import logging
4
+ import shutil
5
+ import tarfile
6
+ from pathlib import Path
7
+
8
+ import click
9
+ import httpx
10
+ import requests
11
+
12
+ from kscale.utils.cli import coro
13
+ from kscale.web.gen.api import SingleArtifactResponse, UploadArtifactResponse
14
+ from kscale.web.utils import get_api_key, get_artifact_dir, get_cache_dir
15
+ from kscale.web.www_client import KScaleWWWClient
16
+
17
+ # Set up logging
18
+ logging.basicConfig(level=logging.INFO)
19
+ logger = logging.getLogger(__name__)
20
+
21
+ ALLOWED_SUFFIXES = {
22
+ ".urdf",
23
+ ".mjcf",
24
+ ".stl",
25
+ ".obj",
26
+ ".dae",
27
+ ".png",
28
+ ".jpg",
29
+ ".jpeg",
30
+ }
31
+
32
+
33
+ async def fetch_urdf_info(artifact_id: str, cache_dir: Path) -> SingleArtifactResponse:
34
+ response_path = cache_dir / "response.json"
35
+ if response_path.exists():
36
+ return SingleArtifactResponse.model_validate_json(response_path.read_text())
37
+ async with KScaleWWWClient() as client:
38
+ response = await client.get_artifact_info(artifact_id)
39
+ response_path.write_text(response.model_dump_json())
40
+ return response
41
+
42
+
43
+ async def download_artifact(artifact_url: str, cache_dir: Path) -> Path:
44
+ filename = cache_dir / Path(artifact_url).name
45
+ headers = {
46
+ "Authorization": f"Bearer {get_api_key()}",
47
+ }
48
+
49
+ if not filename.exists():
50
+ logger.info("Downloading artifact from %s", artifact_url)
51
+
52
+ async with httpx.AsyncClient() as client:
53
+ response = await client.get(artifact_url, headers=headers)
54
+ response.raise_for_status()
55
+ filename.write_bytes(response.content)
56
+ logger.info("Artifact downloaded to %s", filename)
57
+ else:
58
+ logger.info("Artifact already cached at %s", filename)
59
+
60
+ # Extract the .tgz file
61
+ extract_dir = cache_dir / filename.stem
62
+ if not extract_dir.exists():
63
+ logger.info("Extracting %s to %s", filename, extract_dir)
64
+ with tarfile.open(filename, "r:gz") as tar:
65
+ tar.extractall(path=extract_dir)
66
+ else:
67
+ logger.info("Artifact already extracted at %s", extract_dir)
68
+
69
+ return extract_dir
70
+
71
+
72
+ def create_tarball(folder_path: Path, output_filename: str, cache_dir: Path) -> Path:
73
+ tarball_path = cache_dir / output_filename
74
+ with tarfile.open(tarball_path, "w:gz") as tar:
75
+ for file_path in folder_path.rglob("*"):
76
+ if file_path.is_file() and file_path.suffix.lower() in ALLOWED_SUFFIXES:
77
+ tar.add(file_path, arcname=file_path.relative_to(folder_path))
78
+ logger.info("Added %s to tarball", file_path)
79
+ else:
80
+ logger.warning("Skipping %s", file_path)
81
+ logger.info("Created tarball %s", tarball_path)
82
+ return tarball_path
83
+
84
+
85
+ async def download_urdf(artifact_id: str) -> Path:
86
+ cache_dir = get_artifact_dir(artifact_id)
87
+ try:
88
+ urdf_info = await fetch_urdf_info(artifact_id, cache_dir)
89
+ artifact_url = urdf_info.urls.large
90
+ return await download_artifact(artifact_url, cache_dir)
91
+
92
+ except requests.RequestException:
93
+ logger.exception("Failed to fetch URDF info")
94
+ raise
95
+
96
+
97
+ async def show_urdf_info(artifact_id: str) -> None:
98
+ try:
99
+ urdf_info = await fetch_urdf_info(artifact_id, get_artifact_dir(artifact_id))
100
+ logger.info("URDF Artifact ID: %s", urdf_info.artifact_id)
101
+ logger.info("URDF URL: %s", urdf_info.urls.large)
102
+ except requests.RequestException:
103
+ logger.exception("Failed to fetch URDF info")
104
+ raise
105
+
106
+
107
+ async def remove_local_urdf(artifact_id: str) -> None:
108
+ try:
109
+ if artifact_id.lower() == "all":
110
+ cache_dir = get_cache_dir()
111
+ if cache_dir.exists():
112
+ logger.info("Removing all local caches at %s", cache_dir)
113
+ shutil.rmtree(cache_dir)
114
+ else:
115
+ logger.error("No local caches found")
116
+ else:
117
+ artifact_dir = get_artifact_dir(artifact_id)
118
+ if artifact_dir.exists():
119
+ logger.info("Removing local cache at %s", artifact_dir)
120
+ shutil.rmtree(artifact_dir)
121
+ else:
122
+ logger.error("No local cache found for artifact %s", artifact_id)
123
+
124
+ except Exception:
125
+ logger.error("Failed to remove local cache")
126
+ raise
127
+
128
+
129
+ async def upload_urdf(listing_id: str, root_dir: Path) -> UploadArtifactResponse:
130
+ tarball_path = create_tarball(root_dir, "robot.tgz", get_artifact_dir(listing_id))
131
+
132
+ async with KScaleWWWClient() as client:
133
+ response = await client.upload_artifact(listing_id, str(tarball_path))
134
+
135
+ logger.info("Uploaded artifacts: %s", [artifact.artifact_id for artifact in response.artifacts])
136
+ return response
137
+
138
+
139
+ async def upload_urdf_cli(listing_id: str, root_dir: Path) -> UploadArtifactResponse:
140
+ response = await upload_urdf(listing_id, root_dir)
141
+ return response
142
+
143
+
144
+ @click.group()
145
+ def cli() -> None:
146
+ """K-Scale URDF Store CLI tool."""
147
+ pass
148
+
149
+
150
+ @cli.command()
151
+ @click.argument("artifact_id")
152
+ @coro
153
+ async def download(artifact_id: str) -> None:
154
+ """Download a URDF artifact."""
155
+ await download_urdf(artifact_id)
156
+
157
+
158
+ @cli.command()
159
+ @click.argument("artifact_id")
160
+ @coro
161
+ async def info(artifact_id: str) -> None:
162
+ """Show information about a URDF artifact."""
163
+ await show_urdf_info(artifact_id)
164
+
165
+
166
+ @cli.command("remove-local")
167
+ @click.argument("artifact_id")
168
+ @coro
169
+ async def remove_local(artifact_id: str) -> None:
170
+ """Remove local cache of a URDF artifact."""
171
+ await remove_local_urdf(artifact_id)
172
+
173
+
174
+ @cli.command()
175
+ @click.argument("listing_id")
176
+ @click.argument("root_dir", type=click.Path(exists=True, path_type=Path))
177
+ @coro
178
+ async def upload(listing_id: str, root_dir: Path) -> None:
179
+ """Upload a URDF artifact."""
180
+ await upload_urdf_cli(listing_id, root_dir)
181
+
182
+
183
+ if __name__ == "__main__":
184
+ # python -m kscale.web.urdf
185
+ cli()
kscale/web/utils.py ADDED
@@ -0,0 +1,48 @@
1
+ """Utility functions for interacting with the K-Scale WWW API."""
2
+
3
+ import os
4
+ from pathlib import Path
5
+
6
+ from kscale.conf import Settings
7
+
8
+ DEFAULT_UPLOAD_TIMEOUT = 300.0 # 5 minutes
9
+
10
+
11
+ def get_api_root() -> str:
12
+ """Returns the base URL for the K-Scale WWW API.
13
+
14
+ This can be overridden when targetting a different server.
15
+
16
+ Returns:
17
+ The base URL for the K-Scale WWW API.
18
+ """
19
+ return os.getenv("KSCALE_API_ROOT", "https://api.kscale.dev")
20
+
21
+
22
+ def get_api_key() -> str:
23
+ """Returns the API key for the K-Scale WWW API.
24
+
25
+ Returns:
26
+ The API key for the K-Scale WWW API.
27
+ """
28
+ api_key = Settings.load().www.api_key
29
+ if api_key is None:
30
+ api_key = os.getenv("KSCALE_API_KEY")
31
+ if not api_key:
32
+ raise ValueError(
33
+ "API key not found! Get one here and set it as the `KSCALE_API_KEY` environment variable or in your "
34
+ "config file: https://kscale.dev/keys"
35
+ )
36
+ return api_key
37
+
38
+
39
+ def get_cache_dir() -> Path:
40
+ """Returns the cache directory for artifacts."""
41
+ return Path(Settings.load().www.cache_dir).expanduser().resolve()
42
+
43
+
44
+ def get_artifact_dir(artifact_id: str) -> Path:
45
+ """Returns the directory for a specific artifact."""
46
+ cache_dir = get_cache_dir() / artifact_id
47
+ cache_dir.mkdir(parents=True, exist_ok=True)
48
+ return cache_dir
@@ -0,0 +1,134 @@
1
+ """Defines a typed client for the K-Scale WWW API."""
2
+
3
+ import logging
4
+ from pathlib import Path
5
+ from types import TracebackType
6
+ from typing import Any, Dict, Type
7
+ from urllib.parse import urljoin
8
+
9
+ import httpx
10
+ from pydantic import BaseModel
11
+
12
+ from kscale.web.gen.api import (
13
+ BodyAddListingListingsAddPost,
14
+ NewListingResponse,
15
+ SingleArtifactResponse,
16
+ UploadArtifactResponse,
17
+ UploadKRecRequest,
18
+ )
19
+ from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_key, get_api_root
20
+
21
+ logger = logging.getLogger(__name__)
22
+
23
+
24
+ class KScaleWWWClient:
25
+ def __init__(self, base_url: str = get_api_root(), upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT) -> None:
26
+ self.base_url = base_url
27
+ self.upload_timeout = upload_timeout
28
+ self._client: httpx.AsyncClient | None = None
29
+
30
+ @property
31
+ def client(self) -> httpx.AsyncClient:
32
+ if self._client is None:
33
+ self._client = httpx.AsyncClient(
34
+ base_url=self.base_url,
35
+ headers={"Authorization": f"Bearer {get_api_key()}"},
36
+ timeout=httpx.Timeout(30.0),
37
+ )
38
+ return self._client
39
+
40
+ async def _request(
41
+ self,
42
+ method: str,
43
+ endpoint: str,
44
+ *,
45
+ params: Dict[str, Any] | None = None,
46
+ data: BaseModel | None = None,
47
+ files: Dict[str, Any] | None = None,
48
+ ) -> Dict[str, Any]:
49
+ url = urljoin(self.base_url, endpoint)
50
+ kwargs: Dict[str, Any] = {"params": params}
51
+
52
+ if data:
53
+ kwargs["json"] = data.dict(exclude_unset=True)
54
+ if files:
55
+ kwargs["files"] = files
56
+
57
+ response = await self.client.request(method, url, **kwargs)
58
+
59
+ if response.is_error:
60
+ logger.error("Error response from K-Scale: %s", response.text)
61
+ response.raise_for_status()
62
+ return response.json()
63
+
64
+ async def get_artifact_info(self, artifact_id: str) -> SingleArtifactResponse:
65
+ data = await self._request("GET", f"/artifacts/info/{artifact_id}")
66
+ return SingleArtifactResponse(**data)
67
+
68
+ async def upload_artifact(self, listing_id: str, file_path: str) -> UploadArtifactResponse:
69
+ file_name = Path(file_path).name
70
+ with open(file_path, "rb") as f:
71
+ files = {"files": (file_name, f, "application/gzip")}
72
+ data = await self._request("POST", f"/artifacts/upload/{listing_id}", files=files)
73
+ return UploadArtifactResponse(**data)
74
+
75
+ async def create_listing(self, request: BodyAddListingListingsAddPost) -> NewListingResponse:
76
+ data = await self._request("POST", "/listings", data=request)
77
+ return NewListingResponse(**data)
78
+
79
+ async def create_krec(self, request: UploadKRecRequest) -> dict:
80
+ """Create a new K-Rec upload and get the presigned URL."""
81
+ return await self._request(
82
+ "POST",
83
+ "/krecs/upload",
84
+ data=request,
85
+ )
86
+
87
+ async def close(self) -> None:
88
+ if self._client is not None:
89
+ await self._client.aclose()
90
+ self._client = None
91
+
92
+ async def __aenter__(self) -> "KScaleWWWClient":
93
+ return self
94
+
95
+ async def __aexit__(
96
+ self,
97
+ exc_type: Type[BaseException] | None,
98
+ exc_val: BaseException | None,
99
+ exc_tb: TracebackType | None,
100
+ ) -> None:
101
+ await self.close()
102
+
103
+ async def upload_to_presigned_url(self, url: str, file_path: str) -> None:
104
+ """Upload a file using a presigned URL."""
105
+ with open(file_path, "rb") as f:
106
+ async with httpx.AsyncClient(timeout=httpx.Timeout(timeout=self.upload_timeout)) as client:
107
+ response = await client.put(url, content=f.read(), headers={"Content-Type": "application/octet-stream"})
108
+ response.raise_for_status()
109
+
110
+ async def get_presigned_url(self, listing_id: str, file_name: str, checksum: str | None = None) -> dict:
111
+ """Get a presigned URL for uploading an artifact."""
112
+ params = {"filename": file_name}
113
+ if checksum:
114
+ params["checksum"] = checksum
115
+ return await self._request("POST", f"/artifacts/presigned/{listing_id}", params=params)
116
+
117
+ async def get_krec_info(self, krec_id: str) -> dict:
118
+ """Get information about a K-Rec."""
119
+ logger.info("Getting K-Rec info for ID: %s", krec_id)
120
+ try:
121
+ data = await self._request("GET", f"/krecs/download/{krec_id}")
122
+ if not isinstance(data, dict):
123
+ logger.error("Server returned unexpected type: %s", type(data))
124
+ logger.error("Response data: %s", data)
125
+ raise ValueError(f"Server returned {type(data)} instead of dictionary")
126
+
127
+ return {
128
+ "url": data.get("url"),
129
+ "filename": data.get("filename"),
130
+ "checksum": data.get("checksum"),
131
+ }
132
+ except Exception as e:
133
+ logger.error("Failed to get K-Rec info: %s", str(e))
134
+ raise
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2023 Benjamin Bolte
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,55 @@
1
+ Metadata-Version: 2.1
2
+ Name: kscale
3
+ Version: 0.0.13
4
+ Summary: The kscale project
5
+ Home-page: https://github.com/kscalelabs/kscale
6
+ Author: Benjamin Bolte
7
+ Requires-Python: >=3.11
8
+ Description-Content-Type: text/markdown
9
+ License-File: LICENSE
10
+ Requires-Dist: omegaconf
11
+ Requires-Dist: email_validator
12
+ Requires-Dist: httpx
13
+ Requires-Dist: requests
14
+ Requires-Dist: pydantic
15
+ Requires-Dist: click
16
+ Requires-Dist: aiofiles
17
+ Requires-Dist: krec
18
+ Provides-Extra: dev
19
+ Requires-Dist: black; extra == "dev"
20
+ Requires-Dist: darglint; extra == "dev"
21
+ Requires-Dist: mypy; extra == "dev"
22
+ Requires-Dist: pytest; extra == "dev"
23
+ Requires-Dist: ruff; extra == "dev"
24
+ Requires-Dist: datamodel-code-generator; extra == "dev"
25
+
26
+ <p align="center">
27
+ <picture>
28
+ <img alt="K-Scale Open Source Robotics" src="https://media.kscale.dev/kscale-open-source-header.png" style="max-width: 100%;">
29
+ </picture>
30
+ </p>
31
+
32
+ <div align="center">
33
+
34
+ [![License](https://img.shields.io/badge/license-MIT-green)](https://github.com/kscalelabs/ksim/blob/main/LICENSE)
35
+ [![Discord](https://img.shields.io/discord/1224056091017478166)](https://discord.gg/k5mSvCkYQh)
36
+ [![Wiki](https://img.shields.io/badge/wiki-humanoids-black)](https://humanoids.wiki)
37
+ <br />
38
+ [![python](https://img.shields.io/badge/-Python_3.11-blue?logo=python&logoColor=white)](https://github.com/pre-commit/pre-commit)
39
+ [![black](https://img.shields.io/badge/Code%20Style-Black-black.svg?labelColor=gray)](https://black.readthedocs.io/en/stable/)
40
+ [![ruff](https://img.shields.io/badge/Linter-Ruff-red.svg?labelColor=gray)](https://github.com/charliermarsh/ruff)
41
+ <br />
42
+ [![Python Checks](https://github.com/kscalelabs/kscale/actions/workflows/test.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/test.yml)
43
+ [![Publish Python Package](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml/badge.svg)](https://github.com/kscalelabs/kscale/actions/workflows/publish.yml)
44
+
45
+ </div>
46
+
47
+ # K-Scale Command Line Interface
48
+
49
+ 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).
50
+
51
+ ## Installation
52
+
53
+ ```bash
54
+ pip install kscale
55
+ ```
@@ -0,0 +1,31 @@
1
+ kscale/requirements-dev.txt,sha256=WI7-ea4IRJakmqVMN8QKhOsDGrghwtvk03aIsFaNSIw,130
2
+ kscale/requirements.txt,sha256=zj7Rg5vE2Fl3ao66K9jmOlhv5_8SybHvHK5-4xOGZBs,141
3
+ kscale/rust.cpython-313-i386-linux-gnu.so,sha256=yyJkndufdse5C2HeWxcJY1ek0AR1CUGOd3KpkOKcEy8,467040
4
+ kscale/__init__.py,sha256=FOc1X2Sjq58lZ0IMP9e0P4HQ2MM_7yBHk-KbuvApm-0,173
5
+ kscale/api.py,sha256=314tgy4e9fbQWsKPYLyofZyiKfKr9RmvstldXs2jDvY,322
6
+ kscale/conf.py,sha256=iI6eWgzIhzmqqy9EhVh6MbsmNpqCbTlSVdLEYfxyjvI,1416
7
+ kscale/cli.py,sha256=kfu-5dCr0DIEDEnF80tg5md56bEI2GsyA7mEuad1cDw,655
8
+ kscale/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
9
+ kscale/artifacts/plane.obj,sha256=x59-IIrWpLjhotChiqT2Ul6U8s0RcHkaEeUZb4KXL1c,348
10
+ kscale/artifacts/__init__.py,sha256=RK8wdybtCJPgdLLJ8R8-YMi1Ph5ojqAKVJZowHONtgo,232
11
+ kscale/artifacts/plane.urdf,sha256=LCiTk14AyTHjkZ1jvsb0hNaEaJUxDb8Z1JjsgpXu3YM,819
12
+ kscale/web/urdf.py,sha256=rJjQBG3q5HBRxebNAccOqD-rXMS-2fpNdh8k5-6p4p0,5821
13
+ kscale/web/utils.py,sha256=mXfrtH2TuxXkSF01QpcUyC7OHAszQPsiFYGZcHHgTDk,1337
14
+ kscale/web/krec.py,sha256=k9AhBHwLoxRDRUSllgQ88C_xMzwHCY4mHRrBsQFP1EA,5910
15
+ kscale/web/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
16
+ kscale/web/kernels.py,sha256=ddp8IEFSLS9raETQZIphDjxbrj9Jd0gfeQ6Fad0ee-Q,7242
17
+ kscale/web/api.py,sha256=f7_BPwMnByXSBDE1Rzuhr7IIqoIBY2wqDGoqFLQ_0ng,3809
18
+ kscale/web/www_client.py,sha256=Vu_rf2NuJ7fSiqTEb1B71OoTurdNrCtSVKc7K8AIk24,4963
19
+ kscale/web/pybullet.py,sha256=_at_1yf63Fzy7LmXbR-4fvdM9LqPmzYiMniqhKU-BtE,7395
20
+ kscale/web/gen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
21
+ kscale/web/gen/api.py,sha256=JqNit9Q6zUuWz-SrWjcKx3BQIzSmL50r0BT1CfT2hbo,23092
22
+ kscale/utils/api_base.py,sha256=Kk_WtRDdJHmOg6NtHmVxVrcfARSUkhfr29ypLch_pO0,112
23
+ kscale/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
24
+ kscale/utils/cli.py,sha256=JoaY5x5SdUx97KmMM9j5AjRRUqqrTlJ9qVckZptEsYA,827
25
+ kscale/utils/checksum.py,sha256=jt6QmmQND9zrOEnUtOfZpLYROhgto4Gh3OpdUWk4tZA,1093
26
+ kscale-0.0.13.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
27
+ kscale-0.0.13.dist-info/RECORD,,
28
+ kscale-0.0.13.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
29
+ kscale-0.0.13.dist-info/entry_points.txt,sha256=c2pxylFuUJe3aQ_GLEgZrrBW98PPNkiH8Lg0NL_vnrA,42
30
+ kscale-0.0.13.dist-info/METADATA,sha256=-Nmkr2zrj4ydnMtEW2miV_dqgoyY5PJuiirwUrIRsF8,2174
31
+ kscale-0.0.13.dist-info/WHEEL,sha256=G6MKoMIjJiKJ6I6EvKfxsF279_KpeXFhyRGaLGRU4C8,216
@@ -0,0 +1,8 @@
1
+ Wheel-Version: 1.0
2
+ Generator: setuptools (75.6.0)
3
+ Root-Is-Purelib: false
4
+ Tag: cp313-cp313-manylinux_2_5_i686
5
+ Tag: cp313-cp313-manylinux1_i686
6
+ Tag: cp313-cp313-manylinux_2_17_i686
7
+ Tag: cp313-cp313-manylinux2014_i686
8
+
@@ -0,0 +1,2 @@
1
+ [console_scripts]
2
+ kscale = kscale.cli:cli
@@ -0,0 +1 @@
1
+ kscale