|
13 | 13 | from ai.backend.common.resilience.policies.retry import BackoffStrategy, RetryArgs, RetryPolicy |
14 | 14 | from ai.backend.common.resilience.resilience import Resilience |
15 | 15 | from ai.backend.manager.data.app_config_fragment.types import ( |
| 16 | + AppConfigFragmentBulkItemError, |
| 17 | + AppConfigFragmentBulkWriteResult, |
16 | 18 | AppConfigFragmentData, |
17 | 19 | AppConfigFragmentSearchResult, |
18 | 20 | ) |
|
25 | 27 | from ai.backend.manager.models.scopes import SearchScope |
26 | 28 | from ai.backend.manager.repositories.app_config_fragment.creators import ( |
27 | 29 | AppConfigFragmentCreatorSpec, |
| 30 | + GatedAppConfigFragmentCreate, |
28 | 31 | ) |
29 | 32 | from ai.backend.manager.repositories.base import ( |
30 | 33 | BatchQuerier, |
| 34 | + BulkCreator, |
31 | 35 | Creator, |
| 36 | + CreatorSpec, |
32 | 37 | ExistsQuerier, |
33 | 38 | Purger, |
34 | 39 | Querier, |
|
58 | 63 | ) |
59 | 64 |
|
60 | 65 |
|
| 66 | +def _not_found_failures( |
| 67 | + pk_values: Sequence[object], |
| 68 | + succeeded_ids: set[AppConfigFragmentID], |
| 69 | + errored_indices: set[int], |
| 70 | +) -> list[AppConfigFragmentBulkItemError]: |
| 71 | + """Per-item not-found errors for targets that neither succeeded nor errored. |
| 72 | +
|
| 73 | + The generic partial bulk ops silently skip a missing primary key (no row returned, |
| 74 | + no error), so the missing targets are reconstructed from the input here. |
| 75 | + """ |
| 76 | + return [ |
| 77 | + AppConfigFragmentBulkItemError( |
| 78 | + index=index, message=f"App config fragment {pk_value} not found" |
| 79 | + ) |
| 80 | + for index, pk_value in enumerate(pk_values) |
| 81 | + if index not in errored_indices and pk_value not in succeeded_ids |
| 82 | + ] |
| 83 | + |
| 84 | + |
61 | 85 | class AppConfigFragmentDBSource: |
62 | 86 | """Database source for app config fragment operations.""" |
63 | 87 |
|
@@ -112,6 +136,110 @@ async def purge(self, purger: Purger[AppConfigFragmentRow]) -> AppConfigFragment |
112 | 136 | raise AppConfigFragmentNotFound(f"App config fragment {purger.pk_value} not found") |
113 | 137 | return result.row.to_data() |
114 | 138 |
|
| 139 | + @app_config_fragment_db_source_resilience.apply() |
| 140 | + async def bulk_create( |
| 141 | + self, |
| 142 | + items: Sequence[GatedAppConfigFragmentCreate], |
| 143 | + ) -> AppConfigFragmentBulkWriteResult: |
| 144 | + """Create many fragments with partial success. |
| 145 | +
|
| 146 | + Every item's allow-list gate is checked in the same transaction as the writes, |
| 147 | + so the gate checks and the inserts commit atomically. Rejected items are |
| 148 | + reported in ``failed`` with their batch index; the allowed items are inserted |
| 149 | + each in its own savepoint, so a failed insert (e.g. a duplicate natural key) |
| 150 | + is also reported per item while the rest are created. |
| 151 | + """ |
| 152 | + async with self._ops.write_ops() as w: |
| 153 | + allowed: list[tuple[int, CreatorSpec[AppConfigFragmentRow]]] = [] |
| 154 | + failed: list[AppConfigFragmentBulkItemError] = [] |
| 155 | + for index, item in enumerate(items): |
| 156 | + if await w.exists(item.only_if): |
| 157 | + allowed.append((index, item.spec)) |
| 158 | + else: |
| 159 | + failed.append( |
| 160 | + AppConfigFragmentBulkItemError( |
| 161 | + index=index, |
| 162 | + message=( |
| 163 | + f"Writing app config {item.spec.config_name!r} at scope " |
| 164 | + f"{item.spec.scope_type.value!r} is not allowed." |
| 165 | + ), |
| 166 | + ) |
| 167 | + ) |
| 168 | + result = await w.bulk_create_partial( |
| 169 | + BulkCreator(specs=[spec for _, spec in allowed]) |
| 170 | + ) |
| 171 | + failed.extend( |
| 172 | + AppConfigFragmentBulkItemError( |
| 173 | + index=allowed[error.index][0], message=str(error.exception) |
| 174 | + ) |
| 175 | + for error in result.errors |
| 176 | + ) |
| 177 | + return AppConfigFragmentBulkWriteResult( |
| 178 | + succeeded=[row.to_data() for row in result.successes], |
| 179 | + failed=sorted(failed, key=lambda e: e.index), |
| 180 | + ) |
| 181 | + |
| 182 | + @app_config_fragment_db_source_resilience.apply() |
| 183 | + async def bulk_update( |
| 184 | + self, |
| 185 | + updaters: Sequence[Updater[AppConfigFragmentRow]], |
| 186 | + ) -> AppConfigFragmentBulkWriteResult: |
| 187 | + """Update many fragments with partial success. |
| 188 | +
|
| 189 | + No write-gate — see ``update``. Each updater runs in its own savepoint: a |
| 190 | + missing target or a failed update is reported in ``failed`` (with its batch |
| 191 | + index) while the rest are updated. The batch shares one transaction, so the |
| 192 | + successful updates commit together. |
| 193 | + """ |
| 194 | + async with self._ops.write_ops() as w: |
| 195 | + result = await w.bulk_update_partial(updaters) |
| 196 | + succeeded = [row.to_data() for row in result.successes] |
| 197 | + failed = [ |
| 198 | + AppConfigFragmentBulkItemError(index=e.index, message=str(e.exception)) |
| 199 | + for e in result.errors |
| 200 | + ] |
| 201 | + failed.extend( |
| 202 | + _not_found_failures( |
| 203 | + [updater.pk_value for updater in updaters], |
| 204 | + {data.id for data in succeeded}, |
| 205 | + {e.index for e in result.errors}, |
| 206 | + ) |
| 207 | + ) |
| 208 | + return AppConfigFragmentBulkWriteResult( |
| 209 | + succeeded=succeeded, |
| 210 | + failed=sorted(failed, key=lambda e: e.index), |
| 211 | + ) |
| 212 | + |
| 213 | + @app_config_fragment_db_source_resilience.apply() |
| 214 | + async def bulk_purge( |
| 215 | + self, |
| 216 | + purgers: Sequence[Purger[AppConfigFragmentRow]], |
| 217 | + ) -> AppConfigFragmentBulkWriteResult: |
| 218 | + """Purge many fragments with partial success. |
| 219 | +
|
| 220 | + No write-gate — see ``update``. Each purger runs in its own savepoint: a |
| 221 | + missing target or a failed delete is reported in ``failed`` (with its batch |
| 222 | + index) while the rest are purged. |
| 223 | + """ |
| 224 | + async with self._ops.write_ops() as w: |
| 225 | + result = await w.bulk_purge_partial(list(purgers)) |
| 226 | + succeeded = [row.to_data() for row in result.successes] |
| 227 | + failed = [ |
| 228 | + AppConfigFragmentBulkItemError(index=e.index, message=str(e.exception)) |
| 229 | + for e in result.errors |
| 230 | + ] |
| 231 | + failed.extend( |
| 232 | + _not_found_failures( |
| 233 | + [purger.pk_value for purger in purgers], |
| 234 | + {data.id for data in succeeded}, |
| 235 | + {e.index for e in result.errors}, |
| 236 | + ) |
| 237 | + ) |
| 238 | + return AppConfigFragmentBulkWriteResult( |
| 239 | + succeeded=succeeded, |
| 240 | + failed=sorted(failed, key=lambda e: e.index), |
| 241 | + ) |
| 242 | + |
115 | 243 | @app_config_fragment_db_source_resilience.apply() |
116 | 244 | async def admin_search(self, querier: BatchQuerier) -> AppConfigFragmentSearchResult: |
117 | 245 | """Superadmin/internal path: query across all fragments with no scope filter.""" |
|
0 commit comments