main.go

package main

import (
	_ "echo-grpc-triton/docs" // docs is generated by Swag CLI, you have to import it.
	triton "echo-grpc-triton/tritonserver"
	"flag"
	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/swaggo/echo-swagger"
	"log"
	"net/http" // Package http provides HTTP client and server implementations.
)

// Flags contains the information to send requests to Triton inference server.
type Flags struct {
	URL     string
	TIMEOUT int64
}

// parseFlags parses the arguments and initialize the flags.
func parseFlags() Flags {
	var flags = Flags{}
	// https://github.com/NVIDIA/triton-inference-server/tree/master/docs/examples/model_repository/simple
	flag.StringVar(&flags.URL, "u", "localhost:8001", "Inference Server URL. Default: localhost:8001")
	flag.Int64Var(&flags.TIMEOUT, "t", 10, "Timeout. Default: 10 Sec.")
	flag.Parse()
	return flags
}

// @title        Triton API Server with Triton gRPC Client
// @contact.name Curt-Park
// @contact.url  https://github.com/Curt-Park
func main() {
	// Parse the args
	flags := parseFlags()
	log.Println("Flags:", flags)

	// Check the gRPC connection well-established
	client := triton.NewGRPCInferenceServiceAPIClient(flags.URL, flags.TIMEOUT)

	// Create a server with echo
	e := echo.New()
	// Logger middleware logs the information about each HTTP request
	e.Use(middleware.Logger())

	// APIs
	e.GET("/", getHealthCheck)
	e.GET("/liveness", client.GetServerLiveness)
	e.GET("/readiness", client.GetServerReadiness)
	e.GET("/model-metadata", client.GetModelMetadata)
	e.GET("/model-stats", client.GetModelInferStats)
	e.POST("/model-load", client.LoadModel)
	e.POST("/model-unload", client.UnloadModel)
	e.POST("/infer", client.Infer)

	// Swagger
	e.GET("/docs/*", echoSwagger.WrapHandler)
	e.Logger.Fatal(e.Start(":8080"))
}

// @Summary     Healthcheck
// @Description It returns true if the api server is alive.
// @Accept      json
// @Produce     json
// @Success     200 {object} bool "API server's liveness"
// @Router      / [get]
func getHealthCheck(c echo.Context) error {
	return c.JSON(http.StatusOK, true)
}