|
1 | 1 | import argparse |
| 2 | +import asyncio |
2 | 3 | import sys |
3 | 4 | from typing import Optional |
4 | 5 |
|
5 | 6 | from redisvl.cli.utils import add_redis_connection_options, create_redis_url |
6 | 7 | from redisvl.migration import ( |
| 8 | + AsyncMigrationExecutor, |
7 | 9 | MigrationExecutor, |
8 | 10 | MigrationPlanner, |
9 | 11 | MigrationValidator, |
@@ -32,7 +34,7 @@ class Migrate: |
32 | 34 | "\tlist List all available indexes", |
33 | 35 | "\twizard Interactively build a migration plan and schema patch", |
34 | 36 | "\tplan Generate a migration plan for a document-preserving drop/recreate migration", |
35 | | - "\tapply Execute a reviewed drop/recreate migration plan", |
| 37 | + "\tapply Execute a reviewed drop/recreate migration plan (use --async for large migrations)", |
36 | 38 | "\testimate Estimate disk space required for a migration plan (dry-run, no mutations)", |
37 | 39 | "\tvalidate Validate a completed migration plan against the live index", |
38 | 40 | "\n", |
@@ -199,11 +201,17 @@ def apply(self): |
199 | 201 | parser = argparse.ArgumentParser( |
200 | 202 | usage=( |
201 | 203 | "rvl migrate apply --plan <migration_plan.yaml> " |
202 | | - "[--resume <checkpoint.yaml>] " |
| 204 | + "[--async] [--resume <checkpoint.yaml>] " |
203 | 205 | "[--report-out <migration_report.yaml>]" |
204 | 206 | ) |
205 | 207 | ) |
206 | 208 | parser.add_argument("--plan", help="Path to migration_plan.yaml", required=True) |
| 209 | + parser.add_argument( |
| 210 | + "--async", |
| 211 | + dest="use_async", |
| 212 | + help="Use async executor (recommended for large migrations with quantization)", |
| 213 | + action="store_true", |
| 214 | + ) |
207 | 215 | parser.add_argument( |
208 | 216 | "--resume", |
209 | 217 | dest="checkpoint_path", |
@@ -246,9 +254,16 @@ def apply(self): |
246 | 254 | if disk_estimate.has_quantization: |
247 | 255 | print(f"\n{disk_estimate.summary()}\n") |
248 | 256 |
|
249 | | - report = self._apply_sync( |
250 | | - plan, redis_url, args.query_check_file, args.checkpoint_path |
251 | | - ) |
| 257 | + if args.use_async: |
| 258 | + report = asyncio.run( |
| 259 | + self._apply_async( |
| 260 | + plan, redis_url, args.query_check_file, args.checkpoint_path |
| 261 | + ) |
| 262 | + ) |
| 263 | + else: |
| 264 | + report = self._apply_sync( |
| 265 | + plan, redis_url, args.query_check_file, args.checkpoint_path |
| 266 | + ) |
252 | 267 |
|
253 | 268 | write_migration_report(report, args.report_out) |
254 | 269 | if args.benchmark_out: |
@@ -319,6 +334,29 @@ def _apply_sync( |
319 | 334 | self._print_apply_result(report) |
320 | 335 | return report |
321 | 336 |
|
| 337 | + async def _apply_async( |
| 338 | + self, |
| 339 | + plan, |
| 340 | + redis_url: str, |
| 341 | + query_check_file: Optional[str], |
| 342 | + checkpoint_path: Optional[str] = None, |
| 343 | + ): |
| 344 | + """Execute migration asynchronously (non-blocking for large quantization jobs).""" |
| 345 | + executor = AsyncMigrationExecutor() |
| 346 | + |
| 347 | + print(f"\nApplying migration to '{plan.source.index_name}' (async mode)...") |
| 348 | + |
| 349 | + report = await executor.apply( |
| 350 | + plan, |
| 351 | + redis_url=redis_url, |
| 352 | + query_check_file=query_check_file, |
| 353 | + progress_callback=self._make_progress_callback(), |
| 354 | + checkpoint_path=checkpoint_path, |
| 355 | + ) |
| 356 | + |
| 357 | + self._print_apply_result(report) |
| 358 | + return report |
| 359 | + |
322 | 360 | def _print_apply_result(self, report) -> None: |
323 | 361 | """Print the result summary after migration apply.""" |
324 | 362 | if report.result == "succeeded": |
|
0 commit comments