nucliadb 6.5.0.post4408__py3-none-any.whl → 6.5.0.post4413__py3-none-any.whl

nucliadb/train/api/utils.py

@@ -24,7 +24,7 @@ from typing import Optional
 from nucliadb.train.utils import get_shard_manager
 
 
-async def get_kb_partitions(kbid: str, prefix: Optional[str] = None):
+async def get_kb_partitions(kbid: str, prefix: Optional[str] = None) -> list[str]:
     shard_manager = get_shard_manager()
     shards = await shard_manager.get_shards_by_kbid_inner(kbid=kbid)
     valid_shards = []

nucliadb/train/api/v1/shards.py

@@ -26,6 +26,7 @@ from fastapi import HTTPException, Request
 from fastapi.responses import StreamingResponse
 from fastapi_versioning import version
 
+from nucliadb.common.cluster.exceptions import ShardNotFound
 from nucliadb.train.api.utils import get_kb_partitions
 from nucliadb.train.api.v1.router import KB_PREFIX, api
 from nucliadb.train.generator import generate_train_data
@@ -49,7 +50,10 @@ async def object_get_response(
     kbid: str,
     shard: str,
 ) -> StreamingResponse:
-    partitions = await get_kb_partitions(kbid, shard)
+    try:
+        partitions = await get_kb_partitions(kbid, prefix=shard)
+    except ShardNotFound:
+        raise HTTPException(status_code=404, detail=f"No shards found for kb")
     if shard not in partitions:
         raise HTTPException(status_code=404, detail=f"Partition {shard} not found")
     trainset, filter_expression = await get_trainset(request)

nucliadb/train/api/v1/trainset.py

@@ -20,9 +20,10 @@
 
 from typing import Optional
 
-from fastapi import Request
+from fastapi import HTTPException, Request
 from fastapi_versioning import version
 
+from nucliadb.common.cluster.exceptions import ShardNotFound
 from nucliadb.train.api.utils import get_kb_partitions
 from nucliadb.train.api.v1.router import KB_PREFIX, api
 from nucliadb_models.resource import NucliaDBRoles
@@ -57,5 +58,8 @@ async def get_partitions_prefix(request: Request, kbid: str, prefix: str) -> Tra
 
 
 async def get_partitions(kbid: str, prefix: Optional[str] = None) -> TrainSetPartitions:
-    all_keys = await get_kb_partitions(kbid, prefix)
+    try:
+        all_keys = await get_kb_partitions(kbid, prefix)
+    except ShardNotFound:
+        raise HTTPException(status_code=404, detail=f"No shards found for kb")
     return TrainSetPartitions(partitions=all_keys)

nucliadb/train/generator.py

@@ -17,12 +17,14 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
-
 from typing import AsyncIterator, Callable, Optional
 
 from fastapi import HTTPException
+from grpc import StatusCode
+from grpc.aio import AioRpcError
 
 from nucliadb.common.cache import resource_cache
+from nucliadb.train import logger
 from nucliadb.train.generators.field_classifier import (
     field_classification_batch_generator,
 )
@@ -93,7 +95,16 @@ async def generate_train_data(
     # This cache size is an arbitrary number, once we have a metric in place and
     # we analyze memory consumption, we can adjust it with more knoweldge
     with resource_cache(size=settings.resource_cache_size):
-        async for item in batch_generator(kbid, trainset, shard_replica_id, filter_expression):
-            payload = item.SerializeToString()
-            yield len(payload).to_bytes(4, byteorder="big", signed=False)
-            yield payload
+        try:
+            async for item in batch_generator(kbid, trainset, shard_replica_id, filter_expression):
+                payload = item.SerializeToString()
+                yield len(payload).to_bytes(4, byteorder="big", signed=False)
+                yield payload
+        except AioRpcError as exc:
+            if exc.code() == StatusCode.NOT_FOUND:
+                logger.warning(
+                    f"Shard not found in nidx. Halting the stream",
+                    extra={"kbid": kbid, "shard": shard, "shard_replica_id": shard_replica_id},
+                )
+                return
+            raise

{nucliadb-6.5.0.post4408.dist-info → nucliadb-6.5.0.post4413.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.5.0.post4408
+Version: 6.5.0.post4413
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -19,11 +19,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.5.0.post4408
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.5.0.post4408
-Requires-Dist: nucliadb-protos>=6.5.0.post4408
-Requires-Dist: nucliadb-models>=6.5.0.post4408
-Requires-Dist: nidx-protos>=6.5.0.post4408
+Requires-Dist: nucliadb-telemetry[all]>=6.5.0.post4413
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.5.0.post4413
+Requires-Dist: nucliadb-protos>=6.5.0.post4413
+Requires-Dist: nucliadb-models>=6.5.0.post4413
+Requires-Dist: nidx-protos>=6.5.0.post4413
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]

{nucliadb-6.5.0.post4408.dist-info → nucliadb-6.5.0.post4413.dist-info}/RECORD

@@ -303,7 +303,7 @@ nucliadb/tests/config.py,sha256=JN_Jhgj-fwM9_8IeO9pwxr6C1PiwRDrXxm67Y38rU30,2080
 nucliadb/tests/vectors.py,sha256=CcNKx-E8LPpyvRyljbmb-Tn_wST9Juw2CBoogWrKiTk,62843
 nucliadb/train/__init__.py,sha256=NVwe5yULoHXb80itIJT8YJYEz2xbiOPQ7_OMys6XJw8,1301
 nucliadb/train/app.py,sha256=z6xlGVVVaJmZZmLPIVTgkjD-wIz5b0NYlXAQp7hBHYw,2652
-nucliadb/train/generator.py,sha256=fwFYal7VsV0EP7J_g3IOJ-WLpjwqrVo0gEP7vxIlxGs,4152
+nucliadb/train/generator.py,sha256=H8JLkQ23QQVo4CAdg1ZZh_cncPz7COEfaDu1l-h-0hM,4616
 nucliadb/train/lifecycle.py,sha256=3HadM4GRsYb2m-v4jtdr9C-KBEBx8GlrJDArPYi3SWQ,1960
 nucliadb/train/models.py,sha256=BmgmMjDsu_1Ih5JDAqo6whhume90q0ASJcDP9dkMQm8,1198
 nucliadb/train/nodes.py,sha256=6QD9ZnAacPyFCzs1jxJzsL3CkVNyADvsfnPjjrvn-NU,5610
@@ -317,11 +317,11 @@ nucliadb/train/upload.py,sha256=fTjH1KEL-0ogf3LV0T6ODO0QdPGwdZShSUtFUCAcUlA,3256
 nucliadb/train/uploader.py,sha256=xdLGz1ToDue9Q_M8A-_KYkO-V6fWKYOZQ6IGM4FuwWA,6424
 nucliadb/train/utils.py,sha256=OxQ8No19nxOxUhwNYDICNt8n2H-gowkAgu2Vt3Hatzk,3163
 nucliadb/train/api/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/train/api/utils.py,sha256=1E5h1dS3dLppXD4k6qH6jmaY0WYa_ChaNHiTW9y336Q,1307
+nucliadb/train/api/utils.py,sha256=7verIqWRzDVZUsPrc2OQMGSKcd_01aG2lheEc0z5rP8,1320
 nucliadb/train/api/v1/__init__.py,sha256=P4vCIv93r_Cq1WFDDNjy_Wg7zBkzx0S4euXwfPy1LA4,928
 nucliadb/train/api/v1/router.py,sha256=ukdxn5q1oMar6NSPobgJczWsSxLCHw6DYKlb3zwCiSo,910
-nucliadb/train/api/v1/shards.py,sha256=olgjQFKeYSSCbpYhglCRq6Q9WsM9hxOZgk6mOBb4g0o,3367
-nucliadb/train/api/v1/trainset.py,sha256=kpnpDgiMWr1FKHZJgwH7hue5kzilA8-i9X0YHlNeHuU,2113
+nucliadb/train/api/v1/shards.py,sha256=0TLDWdrMwLAD3lb0AmdvpJF_gALxK-dw0ns_lfGSOzY,3553
+nucliadb/train/api/v1/trainset.py,sha256=1Iep1Ze0mmRREHsP3mpJ8cC6sU-_avlqh5ItTRQnxP8,2307
 nucliadb/train/generators/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/train/generators/field_classifier.py,sha256=UcA5snqLNjIHw0VBzXo9ZtSua6o7wBU3tV9_d5qWpRA,3542
 nucliadb/train/generators/field_streaming.py,sha256=p0xu39D5gaSQc-LagKwpgsVaxm2ULTkWZDPi-Ad1lHc,8378
@@ -370,8 +370,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vF0NkFTXiXhXq3bCVXXVV-ED38ECVoUeeYViP8uMqcU,8357
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.5.0.post4408.dist-info/METADATA,sha256=5PRi3_KM32DTF2_BSWmfndESWte8q1aH9fPh0SYJ6aE,4152
-nucliadb-6.5.0.post4408.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.5.0.post4408.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.5.0.post4408.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.5.0.post4408.dist-info/RECORD,,
+nucliadb-6.5.0.post4413.dist-info/METADATA,sha256=B1pDmizFK3zLtNJQUCZAOb4jLeQgnyw3vcryDfJO5y4,4152
+nucliadb-6.5.0.post4413.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.5.0.post4413.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.5.0.post4413.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.5.0.post4413.dist-info/RECORD,,