85 lines
2.1 KiB
Python
85 lines
2.1 KiB
Python
from typing import Annotated
|
|
|
|
import typer
|
|
from google.cloud import aiplatform
|
|
|
|
from rag_eval.config import settings
|
|
|
|
app = typer.Typer()
|
|
|
|
|
|
@app.command()
|
|
def main(
|
|
pipeline_spec_path: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--pipeline-spec-path",
|
|
"-p",
|
|
help="Path to the compiled pipeline YAML file.",
|
|
),
|
|
],
|
|
input_table: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--input-table",
|
|
"-i",
|
|
help="Full BigQuery table name for input (e.g., 'project.dataset.table')",
|
|
),
|
|
],
|
|
output_table: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--output-table",
|
|
"-o",
|
|
help="Full BigQuery table name for output (e.g., 'project.dataset.table')",
|
|
),
|
|
],
|
|
project_id: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--project-id",
|
|
help="Google Cloud project ID.",
|
|
),
|
|
] = settings.project_id,
|
|
location: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--location",
|
|
help="Google Cloud location for the pipeline job.",
|
|
),
|
|
] = settings.location,
|
|
display_name: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--display-name",
|
|
help="Display name for the pipeline job.",
|
|
),
|
|
] = "search-eval-pipeline-job",
|
|
):
|
|
"""Submits a Vertex AI pipeline job."""
|
|
parameter_values = {
|
|
"project_id": project_id,
|
|
"location": location,
|
|
"input_table": input_table,
|
|
"output_table": output_table,
|
|
}
|
|
|
|
job = aiplatform.PipelineJob(
|
|
display_name=display_name,
|
|
template_path=pipeline_spec_path,
|
|
pipeline_root=f"gs://{settings.bucket}/pipeline_root",
|
|
parameter_values=parameter_values,
|
|
project=project_id,
|
|
location=location,
|
|
)
|
|
|
|
print(f"Submitting pipeline job with parameters: {parameter_values}")
|
|
job.submit(
|
|
service_account="sa-cicd-gitlab@bnt-orquestador-cognitivo-dev.iam.gserviceaccount.com"
|
|
)
|
|
print(f"Pipeline job submitted. You can view it at: {job._dashboard_uri()}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app()
|