kscale 0.0.11__py3-none-any.whl → 0.1.1__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
kscale/store/urdf.py DELETED
@@ -1,193 +0,0 @@
1
- """Utility functions for managing artifacts in the K-Scale store."""
2
-
3
- import argparse
4
- import asyncio
5
- import logging
6
- import shutil
7
- import sys
8
- import tarfile
9
- from pathlib import Path
10
- from typing import Literal, Sequence, get_args
11
-
12
- import httpx
13
- import requests
14
-
15
- from kscale.conf import Settings
16
- from kscale.store.client import KScaleStoreClient
17
- from kscale.store.gen.api import SingleArtifactResponse, UploadArtifactResponse
18
- from kscale.store.utils import get_api_key
19
-
20
- # Set up logging
21
- logging.basicConfig(level=logging.INFO)
22
- logger = logging.getLogger(__name__)
23
-
24
- ALLOWED_SUFFIXES = {
25
- ".urdf",
26
- ".mjcf",
27
- ".stl",
28
- ".obj",
29
- ".dae",
30
- ".png",
31
- ".jpg",
32
- ".jpeg",
33
- }
34
-
35
-
36
- def get_cache_dir() -> Path:
37
- return Path(Settings.load().store.cache_dir).expanduser().resolve()
38
-
39
-
40
- def get_artifact_dir(artifact_id: str) -> Path:
41
- cache_dir = get_cache_dir() / artifact_id
42
- cache_dir.mkdir(parents=True, exist_ok=True)
43
- return cache_dir
44
-
45
-
46
- async def fetch_urdf_info(artifact_id: str, cache_dir: Path) -> SingleArtifactResponse:
47
- response_path = cache_dir / "response.json"
48
- if response_path.exists():
49
- return SingleArtifactResponse.model_validate_json(response_path.read_text())
50
- async with KScaleStoreClient() as client:
51
- response = await client.get_artifact_info(artifact_id)
52
- response_path.write_text(response.model_dump_json())
53
- return response
54
-
55
-
56
- async def download_artifact(artifact_url: str, cache_dir: Path) -> Path:
57
- filename = cache_dir / Path(artifact_url).name
58
- headers = {
59
- "Authorization": f"Bearer {get_api_key()}",
60
- }
61
-
62
- if not filename.exists():
63
- logger.info("Downloading artifact from %s", artifact_url)
64
-
65
- async with httpx.AsyncClient() as client:
66
- response = await client.get(artifact_url, headers=headers)
67
- response.raise_for_status()
68
- filename.write_bytes(response.content)
69
- logger.info("Artifact downloaded to %s", filename)
70
- else:
71
- logger.info("Artifact already cached at %s", filename)
72
-
73
- # Extract the .tgz file
74
- extract_dir = cache_dir / filename.stem
75
- if not extract_dir.exists():
76
- logger.info("Extracting %s to %s", filename, extract_dir)
77
- with tarfile.open(filename, "r:gz") as tar:
78
- tar.extractall(path=extract_dir)
79
- else:
80
- logger.info("Artifact already extracted at %s", extract_dir)
81
-
82
- return extract_dir
83
-
84
-
85
- def create_tarball(folder_path: Path, output_filename: str, cache_dir: Path) -> Path:
86
- tarball_path = cache_dir / output_filename
87
- with tarfile.open(tarball_path, "w:gz") as tar:
88
- for file_path in folder_path.rglob("*"):
89
- if file_path.is_file() and file_path.suffix.lower() in ALLOWED_SUFFIXES:
90
- tar.add(file_path, arcname=file_path.relative_to(folder_path))
91
- logger.info("Added %s to tarball", file_path)
92
- else:
93
- logger.warning("Skipping %s", file_path)
94
- logger.info("Created tarball %s", tarball_path)
95
- return tarball_path
96
-
97
-
98
- async def download_urdf(artifact_id: str) -> Path:
99
- cache_dir = get_artifact_dir(artifact_id)
100
- try:
101
- urdf_info = await fetch_urdf_info(artifact_id, cache_dir)
102
- artifact_url = urdf_info.urls.large
103
- return await download_artifact(artifact_url, cache_dir)
104
-
105
- except requests.RequestException:
106
- logger.exception("Failed to fetch URDF info")
107
- raise
108
-
109
-
110
- async def show_urdf_info(artifact_id: str) -> None:
111
- try:
112
- urdf_info = await fetch_urdf_info(artifact_id, get_artifact_dir(artifact_id))
113
- logger.info("URDF Artifact ID: %s", urdf_info.artifact_id)
114
- logger.info("URDF URL: %s", urdf_info.urls.large)
115
- except requests.RequestException:
116
- logger.exception("Failed to fetch URDF info")
117
- raise
118
-
119
-
120
- async def remove_local_urdf(artifact_id: str) -> None:
121
- try:
122
- if artifact_id.lower() == "all":
123
- cache_dir = get_cache_dir()
124
- if cache_dir.exists():
125
- logger.info("Removing all local caches at %s", cache_dir)
126
- shutil.rmtree(cache_dir)
127
- else:
128
- logger.error("No local caches found")
129
- else:
130
- artifact_dir = get_artifact_dir(artifact_id)
131
- if artifact_dir.exists():
132
- logger.info("Removing local cache at %s", artifact_dir)
133
- shutil.rmtree(artifact_dir)
134
- else:
135
- logger.error("No local cache found for artifact %s", artifact_id)
136
-
137
- except Exception:
138
- logger.error("Failed to remove local cache")
139
- raise
140
-
141
-
142
- async def upload_urdf(listing_id: str, root_dir: Path) -> UploadArtifactResponse:
143
- tarball_path = create_tarball(root_dir, "robot.tgz", get_artifact_dir(listing_id))
144
-
145
- async with KScaleStoreClient() as client:
146
- response = await client.upload_artifact(listing_id, str(tarball_path))
147
-
148
- logger.info("Uploaded artifacts: %s", [artifact.artifact_id for artifact in response.artifacts])
149
- return response
150
-
151
-
152
- async def upload_urdf_cli(listing_id: str, args: Sequence[str]) -> UploadArtifactResponse:
153
- parser = argparse.ArgumentParser(description="K-Scale URDF Store", add_help=False)
154
- parser.add_argument("root_dir", type=Path, help="The path to the root directory to upload")
155
- parsed_args = parser.parse_args(args)
156
-
157
- root_dir = parsed_args.root_dir
158
- response = await upload_urdf(listing_id, root_dir)
159
- return response
160
-
161
-
162
- Command = Literal["download", "info", "upload", "remove-local"]
163
-
164
-
165
- async def main(args: Sequence[str] | None = None) -> None:
166
- parser = argparse.ArgumentParser(description="K-Scale URDF Store", add_help=False)
167
- parser.add_argument("command", choices=get_args(Command), help="The command to run")
168
- parser.add_argument("id", help="The ID to use (artifact when downloading, listing when uploading)")
169
- parsed_args, remaining_args = parser.parse_known_args(args)
170
-
171
- command: Command = parsed_args.command
172
- id: str = parsed_args.id
173
-
174
- match command:
175
- case "download":
176
- await download_urdf(id)
177
-
178
- case "info":
179
- await show_urdf_info(id)
180
-
181
- case "remove-local":
182
- await remove_local_urdf(id)
183
-
184
- case "upload":
185
- await upload_urdf_cli(id, remaining_args)
186
-
187
- case _:
188
- logger.error("Invalid command")
189
- sys.exit(1)
190
-
191
-
192
- if __name__ == "__main__":
193
- asyncio.run(main())
kscale/store/utils.py DELETED
@@ -1,33 +0,0 @@
1
- """Utility functions for interacting with the K-Scale Store API."""
2
-
3
- import os
4
-
5
- from kscale.conf import Settings
6
-
7
-
8
- def get_api_root() -> str:
9
- """Returns the base URL for the K-Scale Store API.
10
-
11
- This can be overridden when targetting a different server.
12
-
13
- Returns:
14
- The base URL for the K-Scale Store API.
15
- """
16
- return os.getenv("KSCALE_API_ROOT", "https://api.kscale.dev")
17
-
18
-
19
- def get_api_key() -> str:
20
- """Returns the API key for the K-Scale Store API.
21
-
22
- Returns:
23
- The API key for the K-Scale Store API.
24
- """
25
- api_key = Settings.load().store.api_key
26
- if api_key is None:
27
- api_key = os.getenv("KSCALE_API_KEY")
28
- if not api_key:
29
- raise ValueError(
30
- "API key not found! Get one here and set it as the `KSCALE_API_KEY` environment variable or in your "
31
- "config file: https://kscale.store/keys"
32
- )
33
- return api_key
@@ -1,26 +0,0 @@
1
- kscale/__init__.py,sha256=4m3jAZmGYlh2S1MiGDtND17NHYndlo25f3NsHy6M3RI,178
2
- kscale/api.py,sha256=xBtKj8rgZ400r1Xx9LRY0AzSgIIttoXdejmhHhdVGS0,333
3
- kscale/conf.py,sha256=9fShFaYTbnrm_eiGjmy8ZtC4Q4m6PQkWPyoF3eNyov8,1424
4
- kscale/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
5
- kscale/requirements-dev.txt,sha256=WI7-ea4IRJakmqVMN8QKhOsDGrghwtvk03aIsFaNSIw,130
6
- kscale/requirements.txt,sha256=gxo_niYIHsmyxKxvIOegv45s_lvdnIzI2iFJ2TZzx_U,103
7
- kscale/artifacts/__init__.py,sha256=RK8wdybtCJPgdLLJ8R8-YMi1Ph5ojqAKVJZowHONtgo,232
8
- kscale/artifacts/plane.obj,sha256=x59-IIrWpLjhotChiqT2Ul6U8s0RcHkaEeUZb4KXL1c,348
9
- kscale/artifacts/plane.urdf,sha256=LCiTk14AyTHjkZ1jvsb0hNaEaJUxDb8Z1JjsgpXu3YM,819
10
- kscale/store/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
11
- kscale/store/api.py,sha256=JBp4h6yz_ESSdB-2FpkD_1-hdI8_iEKh9svzkyy3jhs,2386
12
- kscale/store/cli.py,sha256=8ygg_1tZzOOHJotEIgSN9pfumcriPmA31sI_FCFQiTo,859
13
- kscale/store/client.py,sha256=R1IDnf2J4ojAcP8nmUUHfXhcHUt4zP0-mxtVI7MIC5U,2664
14
- kscale/store/pybullet.py,sha256=zoeATQStuRWgmPhku65xjfgvE3Y8ReheUIAkZnDr2C0,7614
15
- kscale/store/urdf.py,sha256=5x8tK2BYv901S_yYWYPWEnHv-3T0ALBQMdDwb70EZFw,6395
16
- kscale/store/utils.py,sha256=euePSkUllLgbJa4RXrCAhB65-BFJOhEImH0ipDs265w,904
17
- kscale/store/gen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
18
- kscale/store/gen/api.py,sha256=GnNTEfFPXVkr0GuZDCCTx3HZ6VD2Yg1ImWpYV-QfmlM,13323
19
- kscale/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
20
- kscale/utils/api_base.py,sha256=Kk_WtRDdJHmOg6NtHmVxVrcfARSUkhfr29ypLch_pO0,112
21
- kscale-0.0.11.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
22
- kscale-0.0.11.dist-info/METADATA,sha256=-BE86Y4RVRIVbvfakjSFbU6QcELc3cFQYofVB01_NkE,2505
23
- kscale-0.0.11.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
24
- kscale-0.0.11.dist-info/entry_points.txt,sha256=PaVs1ivqB0BBdGUsiFkxGUYjGLz05VqagxwRVwi4yV4,54
25
- kscale-0.0.11.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
26
- kscale-0.0.11.dist-info/RECORD,,
@@ -1,2 +0,0 @@
1
- [console_scripts]
2
- kscale = kscale.store.cli:sync_main
File without changes
File without changes