|
20 | 20 | from typing import Any, Dict, List, Optional |
21 | 21 |
|
22 | 22 | import boto3 |
23 | | -from boto3.dynamodb.conditions import Key |
24 | 23 |
|
25 | 24 | logger = logging.getLogger() |
26 | 25 | logger.setLevel(logging.INFO) |
|
39 | 38 | FINETUNING_JOB_PREFIX = "finetuning#" |
40 | 39 | FINETUNING_JOBS_GSI_PK = "finetuning#jobs" |
41 | 40 |
|
42 | | -# Supported base models for fine-tuning |
| 41 | +# Supported base models for fine-tuning (Nova 2.x recommended) |
43 | 42 | SUPPORTED_BASE_MODELS = [ |
44 | | - {"id": "us.amazon.nova-lite-v1:0", "name": "Nova Lite", "provider": "Amazon"}, |
45 | | - {"id": "us.amazon.nova-pro-v1:0", "name": "Nova Pro", "provider": "Amazon"}, |
| 43 | + {"id": "us.amazon.nova-2-lite-v1:0", "name": "Nova 2 Lite", "provider": "Amazon"}, |
| 44 | + {"id": "us.amazon.nova-2-pro-v1:0", "name": "Nova 2 Pro", "provider": "Amazon"}, |
46 | 45 | ] |
47 | 46 |
|
48 | 47 |
|
@@ -95,47 +94,56 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> Any: |
95 | 94 |
|
96 | 95 | def list_finetuning_jobs(arguments: Dict[str, Any]) -> Dict[str, Any]: |
97 | 96 | """List all fine-tuning jobs with pagination. |
98 | | - |
| 97 | +
|
99 | 98 | Uses scan with filter since fine-tuning jobs are relatively few |
100 | 99 | and the GSI1 index may not exist on all deployments. |
101 | 100 | Jobs are stored with GSI1PK/GSI1SK for future GSI support. |
| 101 | +
|
| 102 | + The scan paginates through the table until enough matching items |
| 103 | + are collected or the entire table has been scanned, because |
| 104 | + DynamoDB's ``Limit`` caps items *evaluated* (before filtering), |
| 105 | + not items *returned*. |
102 | 106 | """ |
103 | | - limit = arguments.get("limit", 20) |
| 107 | + limit = arguments.get("limit", 50) |
104 | 108 | next_token = arguments.get("nextToken") |
105 | 109 |
|
106 | 110 | table = dynamodb.Table(TRACKING_TABLE_NAME) |
107 | 111 |
|
108 | | - # Use scan with filter - fine-tuning jobs have PK starting with 'finetuning#' |
109 | 112 | from boto3.dynamodb.conditions import Attr |
110 | | - |
111 | | - scan_params = { |
112 | | - "FilterExpression": Attr("PK").begins_with(FINETUNING_JOB_PREFIX) & Attr("SK").eq("metadata"), |
113 | | - "Limit": limit * 5, # Scan more items since filter reduces results |
114 | | - } |
| 113 | + |
| 114 | + filter_expr = Attr("PK").begins_with(FINETUNING_JOB_PREFIX) & Attr("SK").eq( |
| 115 | + "metadata" |
| 116 | + ) |
| 117 | + |
| 118 | + # Collect all matching items by paginating through the scan. |
| 119 | + # Fine-tuning jobs are few relative to the rest of the table, |
| 120 | + # so a single scan page may not contain any matches. |
| 121 | + items: List[Dict[str, Any]] = [] |
| 122 | + scan_kwargs: Dict[str, Any] = {"FilterExpression": filter_expr} |
115 | 123 |
|
116 | 124 | if next_token: |
117 | | - scan_params["ExclusiveStartKey"] = json.loads(next_token) |
| 125 | + scan_kwargs["ExclusiveStartKey"] = json.loads(next_token) |
118 | 126 |
|
119 | | - response = table.scan(**scan_params) |
| 127 | + while True: |
| 128 | + response = table.scan(**scan_kwargs) |
120 | 129 |
|
121 | | - # Filter and format items |
122 | | - items = [] |
123 | | - for item in response.get("Items", []): |
124 | | - if item.get("PK", "").startswith(FINETUNING_JOB_PREFIX) and item.get("SK") == "metadata": |
| 130 | + for item in response.get("Items", []): |
125 | 131 | items.append(_format_job_for_graphql(item)) |
126 | | - |
| 132 | + |
| 133 | + # Stop if we've scanned the whole table |
| 134 | + if "LastEvaluatedKey" not in response: |
| 135 | + break |
| 136 | + |
| 137 | + # Continue scanning from where we left off |
| 138 | + scan_kwargs["ExclusiveStartKey"] = response["LastEvaluatedKey"] |
| 139 | + |
127 | 140 | # Sort by createdAt descending (most recent first) |
128 | 141 | items.sort(key=lambda x: x.get("createdAt", ""), reverse=True) |
129 | | - |
| 142 | + |
130 | 143 | # Apply limit after sorting |
131 | 144 | items = items[:limit] |
132 | 145 |
|
133 | | - result = {"items": items} |
134 | | - |
135 | | - if "LastEvaluatedKey" in response and len(items) >= limit: |
136 | | - result["nextToken"] = json.dumps(response["LastEvaluatedKey"], cls=DecimalEncoder) |
137 | | - |
138 | | - return result |
| 146 | + return {"items": items} |
139 | 147 |
|
140 | 148 |
|
141 | 149 | def get_finetuning_job(job_id: str) -> Optional[Dict[str, Any]]: |
|
0 commit comments