Computation Control Service

header {
  name: "computations/abc123"
  liquid_legions_v2 {
    description: EXECUTION_PHASE_ONE_INPUT
  }
}

import grpc
from wfa.measurement.system.v1alpha import computation_control_service_pb2
from wfa.measurement.system.v1alpha import computation_control_service_pb2_grpc

def advance_computation_streaming(stub, computation_name, encrypted_data):
    """
    Advance a computation by streaming encrypted data to another duchy.
    
    Args:
        stub: ComputationControl gRPC stub
        computation_name: Resource name of the computation
        encrypted_data: Encrypted sketch data bytes
    """
    def request_iterator():
        # First message: header
        yield computation_control_service_pb2.AdvanceComputationRequest(
            header=computation_control_service_pb2.AdvanceComputationRequest.Header(
                name=computation_name,
                liquid_legions_v2=computation_control_service_pb2.LiquidLegionsV2(
                    description=computation_control_service_pb2.LiquidLegionsV2.EXECUTION_PHASE_ONE_INPUT
                )
            )
        )
        
        # Subsequent messages: data chunks
        chunk_size = 1024 * 1024  # 1MB chunks
        for i in range(0, len(encrypted_data), chunk_size):
            chunk = encrypted_data[i:i + chunk_size]
            yield computation_control_service_pb2.AdvanceComputationRequest(
                body_chunk=computation_control_service_pb2.AdvanceComputationRequest.BodyChunk(
                    partial_data=chunk
                )
            )
    
    # Send streaming request
    response = stub.AdvanceComputation(request_iterator())
    return response

# Usage
with grpc.secure_channel('duchy.example.com:443', credentials) as channel:
    stub = computation_control_service_pb2_grpc.ComputationControlStub(channel)
    advance_computation_streaming(
        stub,
        "computations/abc123",
        encrypted_sketch_data
    )

def get_computation_stage(stub, computation_name, duchy_id):
    """
    Retrieve the current stage of a computation for a duchy.
    
    Args:
        stub: ComputationControl gRPC stub
        computation_name: Computation resource name
        duchy_id: Duchy identifier
    
    Returns:
        ComputationStage message
    """
    stage_name = f"{computation_name}/participants/{duchy_id}/computationStage"
    
    request = computation_control_service_pb2.GetComputationStageRequest(
        name=stage_name
    )
    
    stage = stub.GetComputationStage(request)
    return stage

# Usage
stage = get_computation_stage(stub, "computations/abc123", "duchy-1")
print(f"Current stage: {stage}")