build_tools/python/e2e_test_artifacts/model_artifacts.py

## Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Helpers that generates paths for model artifacts."""

import pathlib
import urllib.parse

from e2e_test_framework.definitions import common_definitions

MODEL_ARTIFACT_PREFIX = "model"
# Archive extensions used to pack models.
ARCHIVE_FILE_EXTENSIONS = [".tar", ".gz"]


def get_model_path(
    model: common_definitions.Model, root_path: pathlib.PurePath = pathlib.PurePath()
) -> pathlib.PurePath:
    """Returns the path of an model artifact file or directory.

    Args:
      model: source model.
      root_path: path of the root artifact directory, on which the returned path
        will base.
    Returns:
      Path of the model artifact.
    """
    model_url = urllib.parse.urlparse(model.source_url)
    # Drop the archive extensions.
    file_exts = pathlib.PurePath(model_url.path).suffixes
    while len(file_exts) > 0 and file_exts[-1] in ARCHIVE_FILE_EXTENSIONS:
        file_exts.pop()
    model_ext = "".join(file_exts)

    # Model path: <root_path>/<model_artifact_prefix>_<model_id>_<model_name><model_ext>
    return root_path / f"{MODEL_ARTIFACT_PREFIX}_{model.id}_{model.name}{model_ext}"