使用 Model Registry 进行模型版本控制

from typing import List

from google.cloud import aiplatform

def upload_new_model_version_using_custom_training_pipeline(
    display_name: str,
    script_path: str,
    container_uri,
    model_serving_container_image_uri: str,
    dataset_id: str,
    replica_count: int,
    machine_type: str,
    accelerator_type: str,
    accelerator_count: int,
    parent_model: str,
    args: List[str],
    model_version_aliases: List[str],
    model_version_description: str,
    is_default_version: bool,
    project: str,
    location: str,
):
    """
    Uploads a new model version using a custom training pipeline.
    Args:
        display_name: The display name of the model version.
        script_path: The path to the Python script that trains the model.
        container_uri: The URI of the container to use for training.
        model_serving_container_image_uri: The URI of the serving container image to use.
        dataset_id: The ID of the dataset to use for training.
        replica_count: The number of replicas to use for training.
        machine_type: The machine type to use for training.
        accelerator_type: The accelerator type to use for training.
        accelerator_count: The number of accelerators to use for training.
        parent_model: The parent resource name of an existing model.
        args: A list of arguments to pass to the training script.
        model_version_aliases: The aliases of the model version to create.
        model_version_description: The description of the model version.
        is_default_version: Whether the model version is the default version.
        project: The project ID.
        location: The region name.
    Returns:
        The new version of the model.
    """
    # Initialize the client.
    aiplatform.init(project=project, location=location)

    # Create the training job.
    # This job will upload a new, non-default version of the my-training-job model
    job = aiplatform.CustomTrainingJob(
        display_name=display_name,
        script_path=script_path,
        container_uri=container_uri,
        model_serving_container_image_uri=model_serving_container_image_uri,
    )

    # Create dataset
    # This examples uses a TabularDataset, but you can use any dataset type.
    dataset = aiplatform.TabularDataset(dataset_id) if dataset_id else None

    # Run the training job.
    model = job.run(
        dataset=dataset,
        args=args,
        replica_count=replica_count,
        machine_type=machine_type,
        accelerator_type=accelerator_type,
        accelerator_count=accelerator_count,
        parent_model=parent_model,
        model_version_aliases=model_version_aliases,
        model_version_description=model_version_description,
        is_default_version=is_default_version,
    )

    return model

from google.cloud import aiplatform

def create_default_model_sample(model_id: str, project: str, location: str):
    """
    Initialize a Model resource to represent an existing model version with alias 'default'.
    Args:
        model_id: The ID of the model to initialize. Parent resource name of the model is also accepted.
        project: The project ID.
        location: The region name.
    Returns:
        Model resource.
    """
    # Initialize the client.
    aiplatform.init(project=project, location=location)

    # Initialize the Model resource with the ID 'model_id'. The parent_name of the Model resource can be also
    # 'projects/<your-project-id>/locations/<your-region>/models/<your-model-id>'
    default_model = aiplatform.Model(model_name=model_id)

    return default_model

from google.cloud import aiplatform

def list_model_versions_sample(model_id: str, project: str, location: str):
    """
    List all model versions of a model.
    Args:
        model_id: The ID of the model to list. Parent resource name of the model is also accepted.
        project: The project ID.
        location: The region name.
    Returns:
        versions: List of model versions.
    """
    # Initialize the client.
    aiplatform.init(project=project, location=location)

    # Initialize the Model Registry resource with the ID 'model_id'.The parent_name of Model resource can be also
    # 'projects/<your-project-id>/locations/<your-region>/models/<your-model-id>'
    model_registry = aiplatform.models.ModelRegistry(model=model_id)

    # List all model versions of the model.
    versions = model_registry.list_versions()

    return versions

from google.cloud import aiplatform

def get_model_version_info_sample(
    model_id: str, version_id: str, project: str, location: str
):
    """
    Get model version info.
    Args:
        model_id: The ID of the model.
        version_id: The version ID of the model version.
        project: The project ID.
        location: The region name.
    Returns:
        VersionInfo resource.
    """

    # Initialize the client.
    aiplatform.init(project=project, location=location)

    # Initialize the Model Registry resource with the ID 'model_id'.The parent_name of Model resource can be also
    # 'projects/<your-project-id>/locations/<your-region>/models/<your-model-id>'
    model_registry = aiplatform.models.ModelRegistry(model=model_id)

    # Get model version info with the version 'version_id'.
    model_version_info = model_registry.get_version_info(version=version_id)

    return model_version_info