explicit.h

Go to the documentation of this file.

/*
    Philote C++ Bindings
 
    Copyright 2022-2025 Christopher A. Lupp
 
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at
 
        http://www.apache.org/licenses/LICENSE-2.0
 
    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
 
    This work has been cleared for public release, distribution unlimited, case
    number: AFRL-2023-5716.
 
    The views expressed are those of the authors and do not reflect the
    official guidance or position of the United States Government, the
    Department of Defense or of the United States Air Force.
 
    Statement from DoD: The Appearance of external hyperlinks does not
    constitute endorsement by the United States Department of Defense (DoD) of
    the linked websites, of the information, products, or services contained
    therein. The DoD does not exercise any editorial, security, or other
    control over the information you may find at these locations.
*/
#pragma once
 
#include <string>
#include <map>
#include <unordered_map>
#include <utility>
 
#include <discipline.h>
#include <variable.h>
#include "discipline_client.h"
 
#include <disciplines.grpc.pb.h>
#include "discipline_server.h"
 
namespace philote
{
    // forward declaration
    class ExplicitDiscipline;
 
    class ExplicitServer : public ExplicitService::Service
    {
    public:
        ExplicitServer() = default;
 
        ~ExplicitServer() noexcept;
 
        void LinkPointers(std::shared_ptr<philote::ExplicitDiscipline> implementation);
 
        void UnlinkPointers();
 
        grpc::Status ComputeFunction(grpc::ServerContext *context,
                                     grpc::ServerReaderWriter<::philote::Array,
                                                             ::philote::Array> *stream) override;
 
        grpc::Status ComputeGradient(grpc::ServerContext *context,
                                     grpc::ServerReaderWriter<::philote::Array,
                                                             ::philote::Array> *stream) override;
 
        // Test helper methods that accept interface pointers for unit testing with mocks
        template<typename StreamType>
        grpc::Status ComputeFunctionImpl(grpc::ServerContext *context, StreamType *stream);
 
        template<typename StreamType>
        grpc::Status ComputeGradientImpl(grpc::ServerContext *context, StreamType *stream);
 
        // Public wrappers for tests
        grpc::Status ComputeFunctionForTesting(grpc::ServerContext *context,
                                               grpc::ServerReaderWriterInterface<::philote::Array,
                                                                                ::philote::Array> *stream) {
            return ComputeFunctionImpl(context, stream);
        }
 
        grpc::Status ComputeGradientForTesting(grpc::ServerContext *context,
                                               grpc::ServerReaderWriterInterface<::philote::Array,
                                                                                ::philote::Array> *stream) {
            return ComputeGradientImpl(context, stream);
        }
 
    private:
        std::shared_ptr<philote::ExplicitDiscipline> implementation_;
    };
 
    class ExplicitDiscipline : public Discipline
    {
    public:
        ExplicitDiscipline();
 
        ~ExplicitDiscipline() noexcept;
 
        void RegisterServices(grpc::ServerBuilder &builder);
 
        virtual void Compute(const philote::Variables &inputs, philote::Variables &outputs);
 
        virtual void ComputePartials(const philote::Variables &inputs,
                                     Partials &partials);
 
    private:
        philote::ExplicitServer explicit_;
        philote::DisciplineServer discipline_server_;
    };
 
    class ExplicitClient : public ::philote::DisciplineClient
    {
    public:
        ExplicitClient() = default;
 
        ~ExplicitClient() noexcept = default;
 
        void ConnectChannel(std::shared_ptr<grpc::ChannelInterface> channel);
 
        Variables ComputeFunction(const Variables &inputs);
 
        Partials ComputeGradient(const Variables &inputs);
 
        void SetStub(std::unique_ptr<ExplicitService::StubInterface> stub)
        {
            stub_ = std::move(stub);
        }
 
    protected:
        std::unique_ptr<ExplicitService::StubInterface> stub_;
    };
}
// Template implementations must be in header
#include <algorithm>
 
namespace philote {
 
template<typename StreamType>
grpc::Status ExplicitServer::ComputeFunctionImpl(grpc::ServerContext *context, StreamType *stream)
{
    if (!implementation_)
    {
        return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Discipline implementation not linked");
    }
 
    if (!stream)
    {
        return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Stream is null");
    }
 
    // Check for cancellation before starting
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before start");
    }
 
    philote::Array array;
 
    // preallocate the inputs based on meta data
    Variables inputs;
    const auto *discipline = static_cast<philote::Discipline *>(implementation_.get());
    if (!discipline)
    {
        return grpc::Status(grpc::StatusCode::INTERNAL, "Failed to cast implementation to Discipline");
    }
 
    for (const auto &var : discipline->var_meta())
    {
        std::string name = var.name();
        if (var.type() == kInput)
            inputs[name] = Variable(var);
    }
 
    // Build O(1) lookup map for variable metadata
    std::unordered_map<std::string, const VariableMetaData*> var_lookup;
    for (const auto &var : discipline->var_meta())
    {
        var_lookup[var.name()] = &var;
    }
 
    while (stream->Read(&array))
    {
        // get variables from the stream message
        const std::string &name = array.name();
 
        // get the variable corresponding to the current message using O(1) lookup
        auto var_it = var_lookup.find(name);
        if (var_it == var_lookup.end())
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Variable not found: " + name);
        }
        const VariableMetaData* var = var_it->second;
 
        // obtain the inputs and discrete inputs from the stream
        if (var->type() == VariableType::kInput)
        {
            try
            {
                // set the variable slice
                inputs[name].AssignChunk(array);
            }
            catch (const std::exception &e)
            {
                return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                              "Failed to assign chunk for variable " + name + ": " + e.what());
            }
        }
        else
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Invalid variable type for input: " + name);
        }
    }
 
    // preallocate outputs
    Variables outputs;
    for (const VariableMetaData &var : discipline->var_meta())
    {
        if (var.type() == kOutput)
            outputs[var.name()] = Variable(var);
    }
 
    // Check for cancellation before expensive computation
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before computation");
    }
 
    // Set context for discipline to check cancellation during compute
    discipline->SetContext(context);
 
    // call the discipline developer-defined Compute function
    try
    {
        implementation_->Compute(inputs, outputs);
    }
    catch (const std::exception &e)
    {
        discipline->ClearContext();
        return grpc::Status(grpc::StatusCode::INTERNAL,
                      "Failed to compute outputs: " + std::string(e.what()));
    }
 
    // Clear context after computation
    discipline->ClearContext();
 
    // Check for cancellation before sending results
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before sending results");
    }
 
    // iterate through continuous outputs
    for (const auto &out : outputs)
    {
        const std::string &name = out.first;
        try
        {
            out.second.Send(name, "", stream, discipline->stream_opts().num_double(), context);
        }
        catch (const std::exception &e)
        {
            return grpc::Status(grpc::StatusCode::INTERNAL,
                          "Failed to send output " + name + ": " + e.what());
        }
    }
 
    return grpc::Status::OK;
}
 
template<typename StreamType>
grpc::Status ExplicitServer::ComputeGradientImpl(grpc::ServerContext *context, StreamType *stream)
{
    if (!implementation_)
    {
        return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Discipline implementation not linked");
    }
 
    if (!stream)
    {
        return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Stream is null");
    }
 
    // Check for cancellation before starting
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before start");
    }
 
    philote::Array array;
 
    // preallocate the inputs based on meta data
    Variables inputs;
    const auto *discipline = static_cast<philote::Discipline *>(implementation_.get());
    if (!discipline)
    {
        return grpc::Status(grpc::StatusCode::INTERNAL, "Failed to cast implementation to Discipline");
    }
 
    for (const auto &var : discipline->var_meta())
    {
        std::string name = var.name();
        if (var.type() == kInput)
            inputs[name] = Variable(var);
    }
 
    // Build O(1) lookup map for variable metadata
    std::unordered_map<std::string, const VariableMetaData*> var_lookup;
    for (const auto &var : discipline->var_meta())
    {
        var_lookup[var.name()] = &var;
    }
 
    while (stream->Read(&array))
    {
        // get variables from the stream message
        const std::string &name = array.name();
 
        // get the variable corresponding to the current message using O(1) lookup
        auto var_it = var_lookup.find(name);
        if (var_it == var_lookup.end())
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Variable not found: " + name);
        }
        const VariableMetaData* var = var_it->second;
 
        // obtain the inputs and discrete inputs from the stream
        if (var->type() == VariableType::kInput)
        {
            try
            {
                // set the variable slice
                inputs[name].AssignChunk(array);
            }
            catch (const std::exception &e)
            {
                return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                              "Failed to assign chunk for variable " + name + ": " + e.what());
            }
        }
        else
        {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Invalid variable type for input: " + name);
        }
    }
 
    // preallocate outputs
    Partials partials;
    for (const PartialsMetaData &par : discipline->partials_meta())
    {
        std::vector<size_t> shape;
        for (const int64_t &dim : par.shape())
            shape.push_back(dim);
 
        partials[std::make_pair(par.name(), par.subname())] = Variable(kOutput, shape);
    }
 
    // Check for cancellation before expensive computation
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before computation");
    }
 
    // Set context for discipline to check cancellation during compute
    discipline->SetContext(context);
 
    // call the discipline developer-defined Compute function
    try
    {
        implementation_->ComputePartials(inputs, partials);
    }
    catch (const std::exception &e)
    {
        discipline->ClearContext();
        return grpc::Status(grpc::StatusCode::INTERNAL,
                      "Failed to compute partials: " + std::string(e.what()));
    }
 
    // Clear context after computation
    discipline->ClearContext();
 
    // Check for cancellation before sending results
    if (context && context->IsCancelled())
    {
        return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled before sending results");
    }
 
    // iterate through continuous outputs
    for (const auto &par : partials)
    {
        const std::string &name = par.first.first;
        const std::string &subname = par.first.second;
        try
        {
            par.second.Send(name, subname, stream, discipline->stream_opts().num_double(), context);
        }
        catch (const std::exception &e)
        {
            return grpc::Status(grpc::StatusCode::INTERNAL,
                          "Failed to send partial " + name + "/" + subname + ": " + e.what());
        }
    }
 
    return grpc::Status::OK;
}
 
} // namespace philote