|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +VGI client that sends parquet data through a function. |
| 4 | +
|
| 5 | +Usage: |
| 6 | + vgi-client --input data.parquet --function echo |
| 7 | + vgi-client --input data.parquet --function sum_all_columns |
| 8 | + vgi-client --input data.parquet --function repeat_inputs --args '[3]' |
| 9 | +""" |
| 10 | + |
| 11 | +import io |
| 12 | +import json |
| 13 | +import subprocess |
| 14 | +import sys |
| 15 | +import traceback |
| 16 | +from typing import Any |
| 17 | + |
| 18 | +import click |
| 19 | +import pyarrow as pa |
| 20 | +import pyarrow.parquet as pq |
| 21 | +import structlog |
| 22 | +from pyarrow import ipc |
| 23 | + |
| 24 | +# Configure structlog to match server output format |
| 25 | +structlog.configure( |
| 26 | + processors=[ |
| 27 | + structlog.processors.add_log_level, |
| 28 | + structlog.processors.TimeStamper(fmt="iso"), |
| 29 | + structlog.dev.ConsoleRenderer(), |
| 30 | + ], |
| 31 | + wrapper_class=structlog.make_filtering_bound_logger(0), # Show all levels |
| 32 | + logger_factory=structlog.PrintLoggerFactory(file=sys.stderr), |
| 33 | +) |
| 34 | + |
| 35 | +log = structlog.get_logger().bind(component="client") |
| 36 | + |
| 37 | + |
| 38 | +def create_init_batch( |
| 39 | + function_name: str, arguments: list[Any], input_schema: pa.Schema |
| 40 | +) -> pa.RecordBatch: |
| 41 | + """Create the initialization batch for the server.""" |
| 42 | + # Create a struct field representing the input schema |
| 43 | + in_type_struct = pa.struct( |
| 44 | + [pa.field(f.name, f.type, f.nullable) for f in input_schema] |
| 45 | + ) |
| 46 | + |
| 47 | + init_schema = pa.schema( |
| 48 | + [ |
| 49 | + pa.field("function_name", pa.string()), |
| 50 | + pa.field("arguments", pa.string()), |
| 51 | + pa.field("in_type", in_type_struct), |
| 52 | + ] |
| 53 | + ) |
| 54 | + |
| 55 | + # Create empty struct value for in_type (actual schema is in the field definition) |
| 56 | + in_type_value = {f.name: None for f in input_schema} |
| 57 | + |
| 58 | + init_batch = pa.RecordBatch.from_pydict( |
| 59 | + { |
| 60 | + "function_name": [function_name], |
| 61 | + "arguments": [json.dumps(arguments)], |
| 62 | + "in_type": [in_type_value], |
| 63 | + }, |
| 64 | + schema=init_schema, |
| 65 | + ) |
| 66 | + |
| 67 | + return init_batch |
| 68 | + |
| 69 | + |
| 70 | +@click.command() |
| 71 | +@click.option( |
| 72 | + "--input", |
| 73 | + "input_file", |
| 74 | + required=True, |
| 75 | + type=click.Path(exists=True), |
| 76 | + help="Path to input parquet file", |
| 77 | +) |
| 78 | +@click.option( |
| 79 | + "--function", |
| 80 | + "function_name", |
| 81 | + required=True, |
| 82 | + type=str, |
| 83 | + help="Name of the function to run (e.g., echo, sum_all_columns, repeat_inputs)", |
| 84 | +) |
| 85 | +@click.option( |
| 86 | + "--args", |
| 87 | + "arguments", |
| 88 | + default="[]", |
| 89 | + type=str, |
| 90 | + help="JSON array of arguments to pass to the function (default: [])", |
| 91 | +) |
| 92 | +@click.option( |
| 93 | + "--server", |
| 94 | + "server_path", |
| 95 | + default="vgi-example-server", |
| 96 | + type=str, |
| 97 | + help="Path to the vgi server", |
| 98 | +) |
| 99 | +def main(input_file: str, function_name: str, arguments: str, server_path: str) -> None: |
| 100 | + """Send parquet data through a VGI function and display results.""" |
| 101 | + # Parse arguments |
| 102 | + try: |
| 103 | + args_list = json.loads(arguments) |
| 104 | + if not isinstance(args_list, list): |
| 105 | + raise click.ClickException("--args must be a JSON array") |
| 106 | + except json.JSONDecodeError as e: |
| 107 | + log.error("invalid_json_arguments", error=str(e)) |
| 108 | + raise click.ClickException(f"Invalid JSON in --args: {e}") from e |
| 109 | + |
| 110 | + # Read input parquet file |
| 111 | + log.info("reading_input", file=input_file) |
| 112 | + table = pq.read_table(input_file) |
| 113 | + input_schema = table.schema |
| 114 | + |
| 115 | + log.info( |
| 116 | + "input_loaded", |
| 117 | + schema=str(input_schema), |
| 118 | + num_rows=table.num_rows, |
| 119 | + num_columns=len(input_schema), |
| 120 | + ) |
| 121 | + |
| 122 | + # Start the server subprocess (stderr goes to screen) |
| 123 | + log.info("starting_server", function=function_name, server_path=server_path) |
| 124 | + proc = subprocess.Popen( |
| 125 | + [sys.executable, server_path], |
| 126 | + stdin=subprocess.PIPE, |
| 127 | + stdout=subprocess.PIPE, |
| 128 | + stderr=None, |
| 129 | + text=False, |
| 130 | + bufsize=0, |
| 131 | + ) |
| 132 | + log.debug("server_started", pid=proc.pid) |
| 133 | + |
| 134 | + assert proc.stdout is not None, "stdout pipe not created" |
| 135 | + stdout_buffered = io.BufferedReader(proc.stdout) # type: ignore[type-var] |
| 136 | + |
| 137 | + try: |
| 138 | + # Send initialization batch |
| 139 | + proc_stdin_sink = pa.PythonFile(proc.stdin) |
| 140 | + log.debug("sending_init_batch", function=function_name, arguments=args_list) |
| 141 | + init_batch = create_init_batch(function_name, args_list, input_schema) |
| 142 | + init_writer = ipc.new_stream(proc_stdin_sink, init_batch.schema) |
| 143 | + init_writer.write_batch(init_batch) |
| 144 | + # If the init_writer is closed here, it closes the pipe to the server |
| 145 | + # which isn't desired, so we leave it open. |
| 146 | + |
| 147 | + log.debug("reading_output_schema") |
| 148 | + msg = ipc.read_message(stdout_buffered) |
| 149 | + assert msg.type == "schema", f"Expected schema message, got {msg.type}" |
| 150 | + output_schema = ipc.read_schema(msg) |
| 151 | + log.info("output_schema_received", schema=str(output_schema)) |
| 152 | + |
| 153 | + # Send data batches |
| 154 | + data_writer = ipc.new_stream(proc_stdin_sink, input_schema) |
| 155 | + batches = table.to_batches() |
| 156 | + log.info("sending_data_batches", num_batches=len(batches)) |
| 157 | + output_reader = None |
| 158 | + for i, input_batch in enumerate(batches): |
| 159 | + while True: |
| 160 | + log.debug("sending_batch", batch_index=i, num_rows=input_batch.num_rows) |
| 161 | + data_writer.write_batch(input_batch) |
| 162 | + log.debug("batch_sent", batch_index=i) |
| 163 | + |
| 164 | + if output_reader is None: |
| 165 | + output_reader = ipc.open_stream(stdout_buffered) |
| 166 | + |
| 167 | + log.debug("attempting_read_output", batch_index=i) |
| 168 | + output_batch, output_metadata = ( |
| 169 | + output_reader.read_next_batch_with_custom_metadata() |
| 170 | + ) |
| 171 | + status = output_metadata.get("status") if output_metadata else None |
| 172 | + |
| 173 | + log.debug( |
| 174 | + "received_output_batch", |
| 175 | + num_rows=output_batch.num_rows, |
| 176 | + status=status, |
| 177 | + batch=output_batch, |
| 178 | + ) |
| 179 | + |
| 180 | + if status == b"HAVE_MORE_OUTPUT": |
| 181 | + # If there are more batches to read for this input, |
| 182 | + # loop here sending the same input batch again, until the |
| 183 | + # server indicates it is done with this input. |
| 184 | + continue |
| 185 | + elif status == b"NEED_MORE_INPUT": |
| 186 | + break |
| 187 | + else: |
| 188 | + log.error("unexpected_status", status=status) |
| 189 | + raise click.ClickException( |
| 190 | + f"Unexpected status from server: {status}" |
| 191 | + ) |
| 192 | + |
| 193 | + empty_input_batch = pa.RecordBatch.from_arrays( |
| 194 | + [pa.array([], type=field.type) for field in input_schema], |
| 195 | + schema=input_schema, |
| 196 | + ) |
| 197 | + |
| 198 | + assert output_reader is not None, "output_reader not initialized" |
| 199 | + while True: |
| 200 | + # Send finalize signal |
| 201 | + log.debug("sending_finalize") |
| 202 | + |
| 203 | + data_writer.write_batch( |
| 204 | + empty_input_batch, custom_metadata={"type": "FINALIZE"} |
| 205 | + ) |
| 206 | + |
| 207 | + output_batch, output_metadata = ( |
| 208 | + output_reader.read_next_batch_with_custom_metadata() |
| 209 | + ) |
| 210 | + status = output_metadata.get("status") if output_metadata else None |
| 211 | + log.debug( |
| 212 | + "received_finalize_batch", |
| 213 | + num_rows=output_batch.num_rows, |
| 214 | + status=status, |
| 215 | + batch=output_batch, |
| 216 | + ) |
| 217 | + |
| 218 | + if status == b"HAVE_MORE_OUTPUT": |
| 219 | + continue |
| 220 | + elif status == b"FINISHED": |
| 221 | + break |
| 222 | + else: |
| 223 | + log.error("unexpected_finalize_status", status=status) |
| 224 | + raise click.ClickException( |
| 225 | + f"Unexpected finalize status from server: {status}" |
| 226 | + ) |
| 227 | + |
| 228 | + data_writer.close() |
| 229 | + proc_stdin_sink.close() |
| 230 | + |
| 231 | + log.info("processing_complete", function=function_name) |
| 232 | + |
| 233 | + except Exception as e: |
| 234 | + log.error("processing_error", error=str(e), traceback=traceback.format_exc()) |
| 235 | + raise click.ClickException(f"Error: {e}") from e |
| 236 | + |
| 237 | + finally: |
| 238 | + proc.wait() |
| 239 | + if proc.returncode != 0: |
| 240 | + log.error("server_exited_with_error", returncode=proc.returncode) |
| 241 | + else: |
| 242 | + log.debug("server_exited", returncode=proc.returncode) |
| 243 | + |
| 244 | + |
| 245 | +if __name__ == "__main__": |
| 246 | + main() |
0 commit comments