feat: Add prediction container URI builder method #805
Conversation
product-auto-label
bot
added
the
api: aiplatform
google-cla
bot
added
the
cla: yes
| framework: str, | ||
| framework_version: str, | ||
| region: str = None, | ||
| with_accelerator: bool = False, |
There was a problem hiding this comment.
I'm unsure about using a flag here, since with PyTorch containers - there are three flavors xla, cpu, and gpu. But as this seems version dependent, the method could handle searching for the correct URI.
There was a problem hiding this comment.
I think we should accept cpu, gpu for now to make this more extendable if we support other accelerators in the future.
| ) | ||
|
|
||
| # All Vertex AI Prediction first-party prediction containers as a string | ||
| _SERVING_CONTAINER_URIS_STR = " ".join(SERVING_CONTAINER_URIS) |
There was a problem hiding this comment.
It seems like it may be simpler to build a map based on the match which will avoid having to maintain ACCELERATOR_TO_URI_REF, the framework constants, and this string in favor of the map.
Something like:
container_map = defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))
for container_uri in SERVING_CONTAINER_URIS:
m = CONTAINER_URI_PATTERN.match(container_uri)
container_map[m['region']][d['framework']][d['accelerator']][d['version']] = container_uriThe TF2 bookkeeping can be done here as well.
Updating with a new container uri or framework will only require adding to the list.
| framework: str, | ||
| framework_version: str, | ||
| region: str = None, | ||
| with_accelerator: bool = False, |
There was a problem hiding this comment.
I think we should accept cpu, gpu for now to make this more extendable if we support other accelerators in the future.
|
|
||
| __all__ = (value_converter,) | ||
| __all__ = ("get_prediction_container_uri" "value_converter",) |
There was a problem hiding this comment.
| __all__ = ("get_prediction_container_uri" "value_converter",) | |
| __all__ = ("get_prediction_container_uri", "value_converter",) |
I'm uncertain of exposing value_converter as well as it's not a module we currently expose or provide documentation/usage on. It mainly used by the enhanced portions of the library.
There was a problem hiding this comment.
Could we potentially move value_converter into aiplatform.utils so that aiplatform.helpers contains only methods intended for external use?
nicain
added
the
type: process
vinnysenthil
changed the title
| f"{DOCS_URI_MESSAGE}" | ||
| ) | ||
|
|
||
| if not URI_MAP[region][framework]: |
There was a problem hiding this comment.
Similarly can use URI_MAP.get(region, {}).get(framework) and the same pattern below.
There was a problem hiding this comment.
Using URI_MAP[region].get(framework) pattern since region is a key that is known to exist at this line.
| raise ValueError( | ||
| f"No serving container for `{framework}` version `{framework_version}` " | ||
| f"with accelerator `{accelerator}` found. Supported versions " | ||
| f"include {', '.join(URI_MAP[region][framework][accelerator].keys())}. {DOCS_URI_MESSAGE}" |
There was a problem hiding this comment.
This code seems to raise error when the user calls the function with version which is already supported by backend, but not yet supported by the installed SDK.
I was asked to mark such case as not an error:
#779 (comment)
Implements a prediction container URI helper as described in this doc and b/197875529