Improve Model Endpoint and Batch Job's Logging UX

api/api/log_api.go

@@ -15,9 +15,9 @@
 package api
 
 import (
-	"bufio"
 	"fmt"
 	"net/http"
+	"time"
 
 	"github.com/gorilla/schema"
 
@@ -42,38 +42,34 @@ func (l *LogController) ReadLog(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	res, err := l.LogService.ReadLog(&query)
-	if err != nil {
-		log.Errorf("Error while retrieving log %v", err)
-		InternalServerError(fmt.Sprintf("Error while retrieving log for container %s: %s", query.Name, err)).WriteTo(w)
-		return
-	}
+	podLogs := make(chan service.PodLog)
+	stopCh := make(chan struct{})
 
 	// send status code and content-type
 	w.Header().Set("Content-Type", "plain/text; charset=UTF-8")
 	w.WriteHeader(http.StatusOK)
 
-	// stream the response body
-	defer res.Close()
-	buff := bufio.NewReader(res)
-	for {
-		line, readErr := buff.ReadString('\n')
-		_, writeErr := w.Write([]byte(line))
-		if writeErr != nil {
-			// connection from caller is closed
-			return
-		}
+	go func() {
+		for podLog := range podLogs {
+			// _, writeErr := w.Write([]byte(podLog.Timestamp.Format(time.RFC3339) + " " + podLog.PodName + "/" + podLog.ContainerName + ": " + podLog.TextPayload + "\n"))
+			_, writeErr := w.Write([]byte(podLog.Timestamp.Format(time.RFC3339) + " " + podLog.TextPayload + "\n"))
+			if writeErr != nil {
+				// connection from caller is closed
+				close(stopCh)
+				return
+			}
 
-		// send the response over network
-		// although it's not guaranteed to reach client if it sits behind proxy
-		flusher, ok := w.(http.Flusher)
-		if ok {
-			flusher.Flush()
+			// send the response over network
+			// although it's not guaranteed to reach client if it sits behind proxy
+			flusher, ok := w.(http.Flusher)
+			if ok {
+				flusher.Flush()
+			}
 		}
+	}()
 
-		if readErr != nil {
-			// unable to read log from container anymore most likely EOF
-			return
-		}
+	if err := l.LogService.StreamLogs(podLogs, stopCh, &query); err != nil {
+		InternalServerError(err.Error())
+		return
 	}
 }

api/cluster/container.go

@@ -56,23 +56,25 @@ func (cf *containerFetcher) GetContainers(namespace string, labelSelector string
 	for _, pod := range podList.Items {
 		for _, c := range pod.Spec.Containers {
 			container := &models.Container{
-				Name:       c.Name,
-				PodName:    pod.Name,
-				Namespace:  pod.Namespace,
-				Cluster:    cf.metadata.ClusterName,
-				GcpProject: cf.metadata.GcpProject,
+				Name:          c.Name,
+				PodName:       pod.Name,
+				ComponentType: models.ComponentType(pod.Name),
+				Namespace:     pod.Namespace,
+				Cluster:       cf.metadata.ClusterName,
+				GcpProject:    cf.metadata.GcpProject,
 			}
 
 			containers = append(containers, container)
 		}
 
 		for _, ic := range pod.Spec.InitContainers {
 			container := &models.Container{
-				Name:       ic.Name,
-				PodName:    pod.Name,
-				Namespace:  pod.Namespace,
-				Cluster:    cf.metadata.ClusterName,
-				GcpProject: cf.metadata.GcpProject,
+				Name:          ic.Name,
+				PodName:       pod.Name,
+				ComponentType: models.ComponentType(pod.Name),
+				Namespace:     pod.Namespace,
+				Cluster:       cf.metadata.ClusterName,
+				GcpProject:    cf.metadata.GcpProject,
 			}
 
 			containers = append(containers, container)

api/imagebuilder/imagebuilder.go

@@ -172,11 +172,12 @@ func (c *imageBuilder) GetContainers(project mlp.Project, model *models.Model, v
 	containers := make([]*models.Container, 0)
 	for _, pod := range pods.Items {
 		containers = append(containers, &models.Container{
-			Name:       containerName,
-			PodName:    pod.Name,
-			Namespace:  pod.Namespace,
-			Cluster:    c.config.ClusterName,
-			GcpProject: c.config.GcpProject,
+			Name:          containerName,
+			PodName:       pod.Name,
+			ComponentType: models.ImageBuilderComponentType,
+			Namespace:     pod.Namespace,
+			Cluster:       c.config.ClusterName,
+			GcpProject:    c.config.GcpProject,
 		})
 	}
 

api/models/container.go

@@ -16,24 +16,46 @@ package models
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/google/uuid"
 )
 
 const (
 	onlineInferenceLabelTemplate = "serving.kubeflow.org/inferenceservice=%s"
 	batchInferenceLabelTemplate  = "prediction-job-id=%s"
+
+	ImageBuilderComponentType     = "image_builder"
+	ModelComponentType            = "model"
+	TransformerComponentType      = "transformer"
+	BatchJobDriverComponentType   = "batch_job_driver"
+	BatchJobExecutorComponentType = "batch_job_executor"
 )
 
 type Container struct {
 	Name              string    `json:"name"`
 	PodName           string    `json:"pod_name"`
+	ComponentType     string    `json:"component_type"`
 	Namespace         string    `json:"namespace"`
 	Cluster           string    `json:"cluster"`
 	GcpProject        string    `json:"gcp_project"`
 	VersionEndpointID uuid.UUID `json:"version_endpoint_id"`
 }
 
+func ComponentType(podName string) string {
+	componentType := ""
+	if strings.Contains(podName, "predictor-default") {
+		componentType = ModelComponentType
+	} else if strings.Contains(podName, "transformer-default") {
+		componentType = TransformerComponentType
+	} else if strings.Contains(podName, "driver") {
+		componentType = BatchJobDriverComponentType
+	} else if strings.Contains(podName, "executor") {
+		componentType = BatchJobDriverComponentType
+	}
+	return componentType
+}
+
 func OnlineInferencePodLabelSelector(modelName string, versionID string) string {
 	serviceName := CreateInferenceServiceName(modelName, versionID)
 	return fmt.Sprintf(onlineInferenceLabelTemplate, serviceName)