Setting the file. One moment.
Search Deployable Models · AWS AI ML · aws/agent-toolkit-for-aws · Skills Docs
Repo No. 14 · AWS AI ML
↖ Back to the coverCustomer Communications Skills
End User Computing Skills
Migration And Modernization Skills
Networking And Content Delivery Skills
Operations Skills · Setting…
Security And Identity Skills
System Table Skills · Querying…
Web And Mobile Development
156 skills · 1155 min
Back to the top of the page 93
Creating Amazon Aurora Db Cluster With Instances
Shieldadvanced
Resilience Program Design
Connecting Lambda To DynamoDB
233
def validate_sdk_versions
— line 233
This file
Number 37.91
Position 91 of 117
Type Python
Size 35 KB
Lines 929 references/model-selection/scripts/ search_deployable_models.py
Python · 929 lines · 35 KB
\
16 --snapshot-file /tmp/sherpa-model-catalog.ABC123.json \
17 --region us-east-2 \
18 "task:text generation" "task:generation-text" "size:<=10b"
19
20 The snapshot is optional for direct developer use. Without one, this script
21 loads the catalog immediately before Search for backward CLI compatibility.
22 The activated skill always supplies the session snapshot so it does not repeat
23 the catalog call after the user confirms filters.
24
25 Dependency: boto3 and botocore must meet the ``MINIMUM_SDK_VERSION`` defined
26 below for native ``Search(Resource="HubContent")`` response support. SageMaker
27 PySDK v3 has no equivalent HubContent Search operation, so this adapter
28 intentionally uses the low-level boto3 client and gates its version before
29 calling Search.
30 """
31
32 import argparse
33 import json
34 import re
35 import sys
36 from collections.abc import Mapping, Sequence
37 from dataclasses import asdict, dataclass
38
39 import boto3
40 import botocore
41 import botocore.exceptions
42
43 from filter_deployable_models import filter_models, sort_models_newest_first
44 from get_deployable_models import (
45 CatalogSnapshotError,
46 get_deployable_models,
47 load_catalog_snapshot,
48 normalize_hub_content,
49 )
50
51 PUBLIC_HUB_NAME = "SageMakerPublicHub"
52 MINIMUM_SDK_VERSION = ( 1 , 43 , 100 )
53 MAX_EXPRESSION_ITEMS = 20
54 MAX_FILTER_VALUE_LENGTH = 1024
55 SEARCH_PAGE_SIZE = 100
56 # Paging key only. Candidate order is restored from OriginalCreationTime.
57 SEARCH_SORT_BY = "HubContentName"
58
59 SUPPORTED_CONSTRAINT_KEYS = (
60 "task" ,
61 "data_type" ,
62 "input_modality" ,
63 "output_modality" ,
64 "license" ,
65 "size" ,
66 "context_window" ,
67 "language" ,
68 "model_type" ,
69 "provider" ,
70 "bedrock" ,
71 "customization" ,
72 "name" ,
73 )
74
75 # field in the normalized catalog, raw Hub keyword prefix, match behavior
76 _KEYWORD_CONSTRAINTS = {
77 "task" : ( "tasks" , "@task:" , "substring" ),
78 "license" : ( "license" , "@license:" , "substring" ),
79 "language" : ( "languages" , "@language:" , "substring" ),
80 "provider" : ( "provider" , "@provider:" , "substring" ),
81 }
82
83 # Constraints served by Search ``Metadata.*`` filters instead of keyword tags.
84 #
85 # The Search index stores every top-level string or string-array field of a
86 # model's HubContentDocument as a metadata entry. A filter names the field as
87 # ``Metadata.<DocumentField>``, and ``Equals``/``In`` compare the stored value
88 # exactly, case included. These fields each have one document source and no
89 # alternate spellings, so exact metadata replaces the derived ``@<tag>:``
90 # keyword one for one. Task, license, language, and provider stay on keywords:
91 # their values are inconsistent (``Apache 2.0`` and ``Apache-2.0``) or capped
92 # by the 50-keyword limit in ways not yet handled.
93 #
94 # Metadata can reach models whose keyword slot ran out. A model with many
95 # languages can exhaust its 50 keywords before ``@model-size:`` or
96 # ``@context-window:`` is written, while the document still carries the field.
97 # The values this adapter sends are still learned from the snapshot, which is
98 # built from ListHubContents and never sees the document. So an overflowed
99 # model is reached only when its raw label is also carried by a model that
100 # kept its keyword. Every overflowed model in the public Hub today shares
101 # ``<1B``, ``1B-10B``, or ``<4K`` with hundreds of tagged models, so all are
102 # reached, and normalize_hub_content reads their size and context window from
103 # the document. A label carried only by overflowed models stays undiscoverable
104 # until values are derived from documents rather than keywords.
105 #
106 # The Hub builds each ``@<tag>:`` keyword by lower-casing the document value,
107 # so the catalog snapshot only knows the lower-cased spelling while metadata
108 # keeps the document's own casing. ``enum`` fields (``ModelTypes``,
109 # ``Capabilities``) are upper-case by schema. Free-text fields arrive from
110 # three producers with three conventions: the Studio manifest lower-cases
111 # (``"text"``), the open-source SDK spec capitalizes (``["Text"]``), and the
112 # proprietary spec upper-cases (``["TEXT"]``). The adapter sends every
113 # convention as exact alternatives so a casing difference cannot silently
114 # empty a result.
115 #
116 # constraint key -> (normalized catalog field, Metadata filter, match, casing)
117 _METADATA_CONSTRAINTS = {
118 "data_type" : ( "data_types" , "Metadata.DataType" , "substring" , "free_text" ),
119 "input_modality" : (
120 "input_modalities" ,
121 "Metadata.InputModalities" ,
122 "substring" ,
123 "free_text" ,
124 ),
125 "output_modality" : (
126 "output_modalities" ,
127 "Metadata.OutputModalities" ,
128 "substring" ,
129 "free_text" ,
130 ),
131 "model_type" : ( "model_type" , "Metadata.ModelTypes" , "exact" , "enum" ),
132 "context_window" : (
133 "context_window" ,
134 "Metadata.ContextWindow" ,
135 "exact" ,
136 "free_text" ,
137 ),
138 }
139
140 # ``size`` is the one constraint whose request value is not a catalog value.
141 # The customer picks a canonical bucket (``<=10b``); the Hub stores raw labels
142 # (``7B``, ``1B-10B``, ``<1B``). The bucket expands to every raw catalog value
143 # that normalize_size_bucket maps into it, and those raw values are what the
144 # metadata filter sends. The bucket mapping stays in Sherpa because raw labels
145 # mix exact counts and ranges and cannot be filtered as a range server-side.
146 _SIZE_METADATA_FILTER = "Metadata.ModelSize"
147
148 # Capability constraints served by ``Metadata.Capabilities``. The document
149 # field is an upper-case enum array; the keyword was its lower-cased copy.
150 _METADATA_CAPABILITIES = {
151 "bedrock" : ( "Metadata.Capabilities" , "BEDROCK_CONSOLE" ),
152 }
153
154 # Capability constraints still served by the lower-cased keyword tag.
155 _CAPABILITY_TAGS = {
156 "customization" : "@capability:customization" ,
157 }
158
159
160 @dataclass ( frozen = True )
161 class OrderingSummary :
162 """Timestamp coverage and provenance for deterministic candidate ordering."""
163
164 field: str
165 timestamp_source: str
166 dated_models: int
167 undated_models: int
168
169
170 @dataclass ( frozen = True )
171 class FallbackReason :
172 """Why public-Hub Search used the complete local catalog fallback."""
173
174 type : str
175 message: str
176
177
178 @dataclass ( frozen = True )
179 class PublicHubSelection :
180 """Typed result for Search-backed public-Hub candidate selection."""
181
182 retrieval_backend: str
183 ordering: OrderingSummary
184 warnings: list[ str ]
185 models: list[ dict ]
186 fallback_reason: FallbackReason | None = None
187
188 def to_output_dict (self):
189 """Return the stable JSON shape, omitting an absent fallback reason."""
190 output = asdict( self )
191 if output[ "fallback_reason" ] is None :
192 del output[ "fallback_reason" ]
193 return output
194
195
196 @dataclass ( frozen = True )
197 class SelectionCatalog :
198 """Catalog data and provenance used for one CLI selection."""
199
200 models: list[ dict ]
201 region_name: str | None
202 source: str
203
204
205 class SearchAdapterError ( Exception ):
206 """Base error for public-Hub Search adapter failures."""
207
208
209 class ConstraintValidationError ( SearchAdapterError ):
210 """A constraint cannot be represented by the public-Hub query contract."""
211
212
213 class UnsupportedSdkError ( SearchAdapterError ):
214 """The installed SDK does not understand HubContent Search records."""
215
216
217 class SearchServiceError ( SearchAdapterError ):
218 """The SageMaker Search request failed."""
219
220
221 class SearchResponseError ( SearchAdapterError ):
222 """The Search response does not contain the expected HubContent shape."""
223
224
225 def _parse_version (version):
226 """Return the first three numeric components of an SDK version."""
227 match = re.match( r " ^\s * (\d + ) \. (\d + ) \. (\d + ) " , str (version))
228 if not match:
229 raise UnsupportedSdkError( f "Cannot parse SDK version: { version !r} " )
230 return tuple ( int (part) for part in match.groups())
231
232
233 def validate_sdk_versions (boto3_version = None , botocore_version = None ):
234 """Require the first boto3/botocore versions with HubContent Search support."""
235 versions = {
236 "boto3" : boto3. __version__ if boto3_version is None else boto3_version,
237 "botocore" : botocore. __version__ if botocore_version is None else botocore_version,
238 }
239 minimum = "." .join( str (part) for part in MINIMUM_SDK_VERSION )
240 for package, version in versions.items():
241 if _parse_version(version) < MINIMUM_SDK_VERSION :
242 raise UnsupportedSdkError(
243 f " { package } { minimum } or newer is required for "
244 f 'Search(Resource="HubContent"); found { version } '
245 )
246
247
248 def normalize_constraints (constraints):
249 """Validate and normalize a mapping of filter keys to value lists."""
250 if not isinstance (constraints, Mapping):
251 raise ConstraintValidationError( "Constraints must be a mapping of keys to value lists" )
252
253 normalized = {}
254 for key, values in constraints.items():
255 if key not in SUPPORTED_CONSTRAINT_KEYS :
256 valid = ", " .join( SUPPORTED_CONSTRAINT_KEYS )
257 raise ConstraintValidationError(
258 f "Unsupported constraint key { key !r} . Supported keys: { valid } "
259 )
260 if isinstance (values, str ) or not isinstance (values, Sequence) or not values:
261 raise ConstraintValidationError(
262 f "Constraint { key !r} must contain at least one value"
263 )
264
265 clean_values = []
266 for value in values:
267 if not isinstance (value, str ) or not value.strip():
268 raise ConstraintValidationError(
269 f "Constraint { key !r} contains an empty or non-string value"
270 )
271 clean_values.append(value.strip().lower())
272 normalized[key] = list ( dict .fromkeys(clean_values))
273
274 return normalized
275
276
277 def parse_constraints (raw_constraints):
278 """Parse repeatable ``key:value`` CLI arguments into normalized constraints."""
279 constraints: dict[ str , list[ str ]] = {}
280 for raw in raw_constraints:
281 if ":" not in raw:
282 raise ConstraintValidationError(
283 f "Malformed constraint { raw !r} ; expected key:value"
284 )
285 key, value = raw.split( ":" , 1 )
286 key = key.strip()
287 constraints.setdefault(key, []).append(value)
288 return normalize_constraints(constraints)
289
290
291 def _catalog_field_values (model, field):
292 """Yield string values from a scalar or list-valued normalized field."""
293 value = model.get(field)
294 if isinstance (value, str ):
295 yield value
296 elif isinstance (value, Sequence):
297 for item in value:
298 if isinstance (item, str ):
299 yield item
300
301
302 def _resolve_catalog_values (key, field, match_mode, requested_values, catalog_models):
303 """Return every current catalog value the customer's request matches.
304
305 Resolution runs against the selection-session snapshot so an unknown value
306 is rejected before any service call and a substring request expands to
307 every exact catalog spelling it covers. Values keep their catalog casing.
308 """
309 resolved: dict[ str , str ] = {}
310
311 for model in catalog_models:
312 for catalog_value in _catalog_field_values(model, field):
313 lowered = catalog_value.lower()
314 if match_mode == "exact" :
315 matches = lowered in requested_values
316 else :
317 matches = any (requested in lowered for requested in requested_values)
318 if matches:
319 resolved.setdefault(catalog_value, catalog_value)
320
321 if not resolved:
322 raise ConstraintValidationError(
323 f "Constraint { key !r} values { requested_values !r} do not match "
324 "any current public-Hub catalog value"
325 )
326 return sorted (resolved.values(), key = str .lower)
327
328
329 def _resolve_keyword_tags (key, requested_values, catalog_models):
330 """Resolve model-selection values to every matching current Hub keyword tag.
331
332 This catalog lookup is intentionally retained for licenses. Public models
333 currently use inconsistent spellings such as ``Apache 2.0`` and
334 ``Apache-2.0``; resolving the customer's substring against the snapshot
335 keeps both exact Search tags until the underlying metadata is standardized.
336 """
337 field, prefix, match_mode = _KEYWORD_CONSTRAINTS [key]
338 values = _resolve_catalog_values(
339 key, field, match_mode, requested_values, catalog_models
340 )
341 return [ f " { prefix }{ value } " for value in values]
342
343
344 def _metadata_spellings (value, casing):
345 """Expand a lower-cased keyword value to the document casings metadata keeps.
346
347 See ``_METADATA_CONSTRAINTS``: keywords are the lower-cased document value,
348 metadata matches the document value exactly. Enum fields have one
349 upper-case spelling; free-text fields may be lower-case, capitalized, or
350 upper-case depending on which Hub producer wrote the document.
351 """
352 lowered = value.lower()
353 if casing == "enum" :
354 return [lowered.upper()]
355 return list ( dict .fromkeys([value, lowered, lowered.capitalize(), lowered.upper()]))
356
357
358 def _resolve_metadata_filter (key, requested_values, catalog_models):
359 """Resolve a constraint to its ``Metadata.*`` filter name and exact values."""
360 field, filter_name, match_mode, casing = _METADATA_CONSTRAINTS [key]
361 values = _resolve_catalog_values(
362 key, field, match_mode, requested_values, catalog_models
363 )
364 spellings: dict[ str , str ] = {}
365 for value in values:
366 for spelling in _metadata_spellings(value, casing):
367 spellings.setdefault(spelling, spelling)
368 return filter_name, list (spellings.values())
369
370
371 def _resolve_size_values (requested_values, catalog_models):
372 """Expand normalized size buckets to the raw ``ModelSize`` values they contain.
373
374 The raw labels come from the catalog snapshot, where they are the Hub's
375 lower-cased keyword copies. The document keeps its own casing (``7B``),
376 so each raw label is expanded to every casing before it is sent.
377 """
378 resolved: dict[ str , str ] = {}
379 for model in catalog_models:
380 normalized_size = model.get( "size" )
381 raw_size = model.get( "size_raw" )
382 if (
383 isinstance (normalized_size, str )
384 and normalized_size.lower() in requested_values
385 and isinstance (raw_size, str )
386 and raw_size
387 ):
388 for spelling in _metadata_spellings(raw_size, "free_text" ):
389 resolved.setdefault(spelling, spelling)
390
391 if not resolved:
392 raise ConstraintValidationError(
393 f "Size values { requested_values !r} do not match any current "
394 "public-Hub raw model-size value"
395 )
396 return sorted (resolved.values(), key = str .lower)
397
398
399 def _resolve_model_names (requested_values, catalog_models):
400 """Resolve case-insensitive substrings to exact current Hub model names."""
401 resolved = set ()
402 for model in catalog_models:
403 name = model.get( "name" )
404 if isinstance (name, str ) and any (
405 requested in name.lower() for requested in requested_values
406 ):
407 resolved.add(name)
408
409 if not resolved:
410 raise ConstraintValidationError(
411 f "Model-name values { requested_values !r} do not match any current "
412 "public-Hub model"
413 )
414 return sorted (resolved, key =lambda name: (name.lower(), name))
415
416
417 def _shared_name_spelling (requested_value, resolved_names):
418 """Return one catalog-cased substring shared by every resolved name."""
419 spellings = set ()
420 for name in resolved_names:
421 start = name.lower().find(requested_value)
422 if start < 0 :
423 return None
424 spellings.add(name[start : start + len (requested_value)])
425 return next ( iter (spellings)) if len (spellings) == 1 else None
426
427
428 def _alternative_filters (name, values):
429 """Pack exact alternatives into safe In filters and literal-comma Equals filters."""
430 search_filters: list[ dict ] = []
431 in_values: list[ str ] = []
432 in_length = 0
433
434 def flush_in_values ():
435 nonlocal in_values, in_length
436 if not in_values:
437 return
438 if len (in_values) == 1 :
439 search_filters.append(
440 { "Name" : name, "Operator" : "Equals" , "Value" : in_values[ 0 ]}
441 )
442 else :
443 search_filters.append(
444 { "Name" : name, "Operator" : "In" , "Value" : "," .join(in_values)}
445 )
446 in_values = []
447 in_length = 0
448
449 for value in values:
450 if "," in value:
451 flush_in_values()
452 search_filters.append(
453 { "Name" : name, "Operator" : "Equals" , "Value" : value}
454 )
455 continue
456
457 added_length = len (value) + ( 1 if in_values else 0 )
458 if in_values and in_length + added_length > MAX_FILTER_VALUE_LENGTH :
459 flush_in_values()
460 added_length = len (value)
461 in_values.append(value)
462 in_length += added_length
463
464 flush_in_values()
465 return search_filters
466
467
468 def _or_expression (search_filters):
469 """Build a nested OR tree whose individual lists stay within service limits."""
470 if len (search_filters) <= MAX_EXPRESSION_ITEMS :
471 return { "Operator" : "Or" , "Filters" : search_filters}
472
473 level = [
474 { "Operator" : "Or" , "Filters" : search_filters[index : index + MAX_EXPRESSION_ITEMS ]}
475 for index in range ( 0 , len (search_filters), MAX_EXPRESSION_ITEMS )
476 ]
477 while len (level) > MAX_EXPRESSION_ITEMS :
478 level = [
479 { "Operator" : "Or" , "SubExpressions" : level[index : index + MAX_EXPRESSION_ITEMS ]}
480 for index in range ( 0 , len (level), MAX_EXPRESSION_ITEMS )
481 ]
482 return { "Operator" : "Or" , "SubExpressions" : level}
483
484
485 def _append_exact_alternatives (filters, subexpressions, name, values):
486 """Add exact alternatives without exceeding filter-list or value limits."""
487 if len (values) == 1 :
488 filters.append({ "Name" : name, "Operator" : "Equals" , "Value" : values[ 0 ]})
489 return
490
491 alternative_filters = _alternative_filters(name, values)
492 if len (alternative_filters) == 1 :
493 filters.append(alternative_filters[ 0 ])
494 return
495 subexpressions.append(_or_expression(alternative_filters))
496
497
498 def validate_search_expression (expression, path = "SearchExpression" ):
499 """Validate generated expression-list and filter-value service limits."""
500 if not isinstance (expression, Mapping):
501 raise ConstraintValidationError( f " { path } must be an object" )
502
503 has_condition = False
504 for collection_name in ( "Filters" , "NestedFilters" , "SubExpressions" ):
505 if collection_name not in expression:
506 continue
507 values = expression[collection_name]
508 if not isinstance (values, list ) or not values:
509 raise ConstraintValidationError( f " { path } . { collection_name } must be a non-empty list" )
510 has_condition = True
511 if len (values) > MAX_EXPRESSION_ITEMS :
512 raise ConstraintValidationError(
513 f " { path } . { collection_name } has { len (values) } elements; "
514 f "the maximum is { MAX_EXPRESSION_ITEMS } "
515 )
516
517 for index, search_filter in enumerate (expression.get( "Filters" , [])):
518 if not isinstance (search_filter, Mapping):
519 raise ConstraintValidationError( f " { path } .Filters[ { index } ] must be an object" )
520 value = search_filter.get( "Value" )
521 if not isinstance (value, str ) or not value:
522 raise ConstraintValidationError(
523 f " { path } .Filters[ { index } ].Value must be a non-empty string"
524 )
525 if len (value) > MAX_FILTER_VALUE_LENGTH :
526 raise ConstraintValidationError(
527 f " { path } .Filters[ { index } ].Value exceeds "
528 f " { MAX_FILTER_VALUE_LENGTH } characters"
529 )
530
531 for index, subexpression in enumerate (expression.get( "SubExpressions" , [])):
532 validate_search_expression(subexpression, f " { path } .SubExpressions[ { index } ]" )
533
534 for index, nested_filter in enumerate (expression.get( "NestedFilters" , [])):
535 if not isinstance (nested_filter, Mapping):
536 raise ConstraintValidationError(
537 f " { path } .NestedFilters[ { index } ] must be an object"
538 )
539 nested_expression = { "Filters" : nested_filter.get( "Filters" , [])}
540 validate_search_expression(
541 nested_expression, f " { path } .NestedFilters[ { index } ]"
542 )
543
544 if not has_condition:
545 raise ConstraintValidationError( f " { path } must contain at least one condition" )
546
547
548 def build_search_expression (hub_name, constraints, catalog_models):
549 """Translate confirmed model-selection constraints into a SearchExpression."""
550 if hub_name != PUBLIC_HUB_NAME :
551 raise ConstraintValidationError(
552 "The Search adapter supports SageMakerPublicHub only; "
553 "private Hubs must continue using ListHubContents"
554 )
555
556 normalized_constraints = normalize_constraints(constraints)
557 filters = [
558 { "Name" : "HubName" , "Operator" : "Equals" , "Value" : PUBLIC_HUB_NAME },
559 { "Name" : "HubContentType" , "Operator" : "Equals" , "Value" : "Model" },
560 ]
561 subexpressions: list[ dict ] = []
562
563 for key in SUPPORTED_CONSTRAINT_KEYS :
564 if key not in normalized_constraints:
565 continue
566 values = normalized_constraints[key]
567
568 if key in _METADATA_CONSTRAINTS :
569 filter_name, spellings = _resolve_metadata_filter(
570 key, values, catalog_models
571 )
572 _append_exact_alternatives(filters, subexpressions, filter_name, spellings)
573 elif key in _KEYWORD_CONSTRAINTS :
574 tags = _resolve_keyword_tags(key, values, catalog_models)
575 _append_exact_alternatives(
576 filters, subexpressions, "HubContentSearchKeywords" , tags
577 )
578 elif key == "size" :
579 raw_sizes = _resolve_size_values(values, catalog_models)
580 _append_exact_alternatives(
581 filters, subexpressions, _SIZE_METADATA_FILTER , raw_sizes
582 )
583 elif key in _METADATA_CAPABILITIES or key in _CAPABILITY_TAGS :
584 # Capabilities are positive exact values. Absence-based false
585 # matching is not enabled until its Search semantics are verified.
586 if values != [ "true" ]:
587 raise ConstraintValidationError(
588 f "The Search adapter supports { key } :true only; "
589 "absence-based false matching is not yet verified"
590 )
591 if key in _METADATA_CAPABILITIES :
592 filter_name, capability = _METADATA_CAPABILITIES [key]
593 filters.append(
594 { "Name" : filter_name, "Operator" : "Equals" , "Value" : capability}
595 )
596 else :
597 _append_exact_alternatives(
598 filters,
599 subexpressions,
600 "HubContentSearchKeywords" ,
601 [ _CAPABILITY_TAGS [key]],
602 )
603 elif key == "name" :
604 names = _resolve_model_names(values, catalog_models)
605 spelling = (
606 _shared_name_spelling(values[ 0 ], names) if len (values) == 1 else None
607 )
608 if spelling is not None :
609 filters.append(
610 {
611 "Name" : "HubContentName" ,
612 "Operator" : "Contains" ,
613 "Value" : spelling,
614 }
615 )
616 else :
617 # Several requested substrings or inconsistent catalog casing
618 # cannot share the service's single case-sensitive Contains.
619 # Exact names preserve local case-insensitive matching.
620 _append_exact_alternatives(
621 filters, subexpressions, "HubContentName" , names
622 )
623
624 expression = { "Operator" : "And" , "Filters" : filters}
625 if subexpressions:
626 expression[ "SubExpressions" ] = subexpressions
627 validate_search_expression(expression)
628 return expression
629
630
631 def _search_error (error):
632 """Convert a ClientError into a stable adapter error."""
633 code = error.response.get( "Error" , {}).get( "Code" , "Unknown" )
634 message = error.response.get( "Error" , {}).get( "Message" , str (error))
635 return SearchServiceError( f "SageMaker Search failed ( { code } ): { message } " )
636
637
638 def search_deployable_models (
639 hub_name,
640 constraints,
641 catalog_models,
642 region_name = None ,
643 sm_client = None ,
644 check_sdk = True ,
645 ):
646 """Search the public Hub and return normalized candidate models."""
647 expression = build_search_expression(hub_name, constraints, catalog_models)
648 if check_sdk:
649 validate_sdk_versions()
650 try :
651 client = sm_client or boto3.client( "sagemaker" , region_name = region_name)
652 except botocore.exceptions.BotoCoreError as error:
653 raise SearchServiceError( f "Could not create SageMaker client: { error } " ) from error
654
655 models = []
656 next_token = None
657 seen_tokens = set ()
658 seen_names = set ()
659
660 while True :
661 request = {
662 "Resource" : "HubContent" ,
663 "SearchExpression" : expression,
664 "MaxResults" : SEARCH_PAGE_SIZE ,
665 # Pages of an unsorted Search are not stable: identical repeated
666 # runs of one multi-page query return the same record count but
667 # different unique names, with some models repeated and others
668 # skipped. A sort key makes the page boundaries deterministic. The
669 # name is used only for paging; restore_catalog_order() applies the
670 # customer-visible OriginalCreationTime order afterwards.
671 "SortBy" : SEARCH_SORT_BY ,
672 "SortOrder" : "Ascending" ,
673 }
674 if next_token:
675 request[ "NextToken" ] = next_token
676
677 try :
678 # SageMaker PySDK v3 does not expose HubContent Search. Use the
679 # version-gated boto3 operation so HubContent records are parsed
680 # instead of silently discarded by an older client model.
681 response = client.search( ** request)
682 except botocore.exceptions.ClientError as error:
683 raise _search_error(error) from error
684 except botocore.exceptions.BotoCoreError as error:
685 raise SearchServiceError( f "SageMaker Search failed: { error } " ) from error
686
687 records = response.get( "Results" , [])
688 if not isinstance (records, list ):
689 raise SearchResponseError( "Search response Results must be a list" )
690 for index, record in enumerate (records):
691 if not isinstance (record, Mapping) or not isinstance (
692 record.get( "HubContent" ), Mapping
693 ):
694 raise SearchResponseError(
695 f "Search result { index } does not contain a HubContent object"
696 )
697 model = normalize_hub_content(record[ "HubContent" ])
698 # A model repeated across pages means paging was not stable; keep
699 # the candidate set a set rather than double-listing it.
700 name = model.get( "name" )
701 if isinstance (name, str ):
702 if name in seen_names:
703 continue
704 seen_names.add(name)
705 models.append(model)
706
707 next_token = response.get( "NextToken" )
708 if next_token is None or next_token == "" :
709 break
710 if not isinstance (next_token, str ):
711 raise SearchResponseError( "Search response NextToken must be a string" )
712 if next_token in seen_tokens:
713 raise SearchResponseError( "Search returned a repeated NextToken" )
714 seen_tokens.add(next_token)
715
716 return models
717
718
719 def _ordering_summary (models, timestamp_source):
720 """Describe timestamp coverage so missing List-only dates stay visible."""
721 dated_models = sum (
722 model.get( "original_creation_time" ) is not None for model in models
723 )
724 return OrderingSummary(
725 field = "OriginalCreationTime" ,
726 timestamp_source = timestamp_source,
727 dated_models = dated_models,
728 undated_models = len (models) - dated_models,
729 )
730
731
732 def restore_catalog_order (
733 search_models,
734 catalog_models,
735 timestamp_source = "selection_session_snapshot" ,
736 ):
737 """Attach catalog dates to Search matches and apply current ordering.
738
739 Search ``CreationTime`` identifies the indexed version and can move when
740 ``latestSupported`` changes. This join deliberately removes any normalized
741 Search timestamp and uses only ListHubContents ``OriginalCreationTime``.
742 A Search-only model stays undated and sorts after dated models by name.
743 """
744 catalog_dates = {
745 model.get( "name" ): model.get( "original_creation_time" )
746 for model in catalog_models
747 if isinstance (model.get( "name" ), str )
748 and isinstance (model.get( "original_creation_time" ), str )
749 }
750 restored = []
751 for model in search_models:
752 candidate = dict (model)
753 candidate.pop( "original_creation_time" , None )
754 catalog_date = catalog_dates.get(candidate.get( "name" ))
755 if catalog_date is not None :
756 candidate[ "original_creation_time" ] = catalog_date
757 restored.append(candidate)
758
759 ordered = sort_models_newest_first(restored)
760 return ordered, _ordering_summary(ordered, timestamp_source)
761
762
763 def select_public_hub_models (
764 hub_name,
765 constraints,
766 catalog_models,
767 region_name = None ,
768 sm_client = None ,
769 catalog_source = "snapshot" ,
770 ):
771 """Use Search for candidates, with a visible complete local fallback.
772
773 A successful empty Search response is authoritative and does not fall back.
774 Only explicit SDK, service, transport, or response failures use the complete
775 ListHubContents snapshot. Constraint and snapshot validation errors remain
776 visible because falling back would hide a query or workflow defect.
777 """
778 if catalog_source == "snapshot" :
779 search_timestamp_source = "selection_session_snapshot"
780 fallback_timestamp_source = "list_hub_contents_snapshot"
781 elif catalog_source == "live_list" :
782 search_timestamp_source = fallback_timestamp_source = "live_list"
783 else :
784 raise ConstraintValidationError(
785 f "Unsupported catalog source: { catalog_source !r} "
786 )
787
788 try :
789 search_models = search_deployable_models(
790 hub_name,
791 constraints,
792 catalog_models,
793 region_name = region_name,
794 sm_client = sm_client,
795 )
796 except (UnsupportedSdkError, SearchServiceError, SearchResponseError) as error:
797 fallback_models = filter_models(catalog_models, constraints)
798 return PublicHubSelection(
799 retrieval_backend = "list_fallback" ,
800 fallback_reason = FallbackReason(
801 type = type (error). __name__ ,
802 message = str (error),
803 ),
804 ordering = _ordering_summary(
805 fallback_models, fallback_timestamp_source
806 ),
807 warnings = [
808 "SageMaker Search was unavailable; used the complete "
809 "ListHubContents snapshot and local filtering instead."
810 ],
811 models = fallback_models,
812 )
813
814 ordered_models, ordering = restore_catalog_order(
815 search_models,
816 catalog_models,
817 timestamp_source = search_timestamp_source,
818 )
819 warnings = []
820 if ordering.undated_models:
821 warnings.append(
822 f " { ordering.undated_models } of { len (ordered_models) } Search results "
823 "had no cached OriginalCreationTime and were ordered after dated models."
824 )
825 return PublicHubSelection(
826 retrieval_backend = "search" ,
827 ordering = ordering,
828 warnings = warnings,
829 models = ordered_models,
830 )
831
832
833 def _parser ():
834 parser = argparse.ArgumentParser(
835 description = "Search SageMakerPublicHub models with confirmed exact constraints."
836 )
837 parser.add_argument( "hub_name" , help = "Must be SageMakerPublicHub" )
838 parser.add_argument( "constraints" , nargs = "*" , help = "Repeatable key:value constraints" )
839 parser.add_argument( "--region" , dest = "region_name" , help = "AWS region override" )
840 parser.add_argument(
841 "--snapshot-file" ,
842 help = "Selection-session catalog snapshot created during --list-values" ,
843 )
844 return parser
845
846
847 def parse_cli_args (argv = None ):
848 """Parse options and repeatable constraints in any documented order."""
849 parser = _parser()
850 # Python 3.10 does not reliably resume a ``nargs='*'`` positional after an
851 # optional argument with parse_args(). parse_intermixed_args() supports the
852 # documented option/constraint order on every runtime we ship.
853 return parser.parse_intermixed_args(argv)
854
855
856 def _load_catalog_for_selection (args):
857 """Load a session snapshot, or preserve direct CLI compatibility."""
858 if args.snapshot_file:
859 snapshot = load_catalog_snapshot(
860 args.snapshot_file,
861 expected_hub = args.hub_name,
862 expected_region = args.region_name,
863 )
864 return SelectionCatalog(
865 models = snapshot[ "catalog_models" ],
866 region_name = snapshot[ "region" ],
867 source = "snapshot" ,
868 )
869
870 try :
871 models = get_deployable_models(args.hub_name, args.region_name)
872 except botocore.exceptions.ClientError as error:
873 code = error.response.get( "Error" , {}).get( "Code" , "Unknown" )
874 message = error.response.get( "Error" , {}).get( "Message" , str (error))
875 raise SearchServiceError(
876 f "Public-Hub catalog discovery failed ( { code } ): { message } "
877 ) from error
878 except botocore.exceptions.BotoCoreError as error:
879 raise SearchServiceError(
880 f "Public-Hub catalog discovery failed: { error } "
881 ) from error
882 return SelectionCatalog(
883 models = models,
884 region_name = args.region_name,
885 source = "live_list" ,
886 )
887
888
889 def main (argv = None ):
890 args = parse_cli_args(argv)
891 try :
892 constraints = parse_constraints(args.constraints)
893 if args.hub_name != PUBLIC_HUB_NAME :
894 raise ConstraintValidationError(
895 "The public-Hub Search adapter supports SageMakerPublicHub only; "
896 "private Hubs must continue using ListHubContents"
897 )
898
899 catalog = _load_catalog_for_selection(args)
900 selection = select_public_hub_models(
901 args.hub_name,
902 constraints,
903 catalog.models,
904 region_name = catalog.region_name,
905 catalog_source = catalog.source,
906 )
907 selection_output = selection.to_output_dict()
908 print (
909 json.dumps(
910 {
911 "total_models" : len (catalog.models),
912 "filters_applied" : constraints,
913 "matched" : len (selection.models),
914 "catalog_source" : catalog.source,
915 ** selection_output,
916 }
917 )
918 )
919 return 0
920 except (SearchAdapterError, CatalogSnapshotError) as error:
921 print ( f "Error: { error } " , file = sys.stderr)
922 return 1
923 except Exception as error:
924 print ( f "Error: { error } " , file = sys.stderr)
925 return 1
926
927
928 if __name__ == "__main__" :
929 raise SystemExit (main())