feat: fetch remote FirecREST logs into sessions (#1136)

* feat: fetch remote logs into sessions

* refactor: use metadata from api

* refactor: use pointers from controller

* fix: offset for partial lines

* fix: stream duplication when stderr is the same as stdout or null

Author: SalimKayal

internal/remote/firecrest/controller.go

@@ -67,6 +67,16 @@ type FirecrestRemoteSessionController struct {
 
 	// fakeStart if true, do not start the remote session and print debug information
 	fakeStart bool
+
+	// stdoutPath and stderrPath are the paths to the Slurm stdout/stderr files on the cluster filesystem.
+	stdoutPath string
+	stderrPath string
+	// stdoutOffset and stderrOffset track how many bytes have been consumed from each log file.
+	stdoutOffset int
+	stderrOffset int
+	// stdoutBuf and stderrBuf hold partial lines between fetches.
+	stdoutBuf bytes.Buffer
+	stderrBuf bytes.Buffer
 }
 
 func NewFirecrestRemoteSessionController(cfg config.RemoteSessionControllerConfig) (c *FirecrestRemoteSessionController, err error) {
@@ -133,7 +143,7 @@ func (c *FirecrestRemoteSessionController) Start(ctx context.Context) error {
 	// Start a go routine to update the session status
 	go c.periodicSessionStatus(ctx)
 
-	if err := c.recoverJobID(); err != nil {
+	if err := c.recoverState(); err != nil {
 		return err
 	}
 	// We recovered an existing job ID, do nothing
@@ -327,8 +337,38 @@ func (c *FirecrestRemoteSessionController) Start(ctx context.Context) error {
 
 	slog.Info("submitted job", "jobID", c.jobID)
 
-	// Save the job ID for recovery
-	if err := c.saveJobID(); err != nil {
+	// After submission, determine the log file paths from the job metadata.
+	c.stdoutPath = path.Join(sessionPath, fmt.Sprintf("slurm-%s.out", c.jobID))
+	c.stderrPath = path.Join(sessionPath, fmt.Sprintf("slurm-%s.err", c.jobID))
+
+	metaRes, err := c.client.GetJobMetadataComputeSystemNameJobsJobIdMetadataGetWithResponse(startCtx, c.systemName, c.jobID)
+	if err != nil {
+		slog.Warn("could not get job metadata, falling back to default log paths", "error", err)
+	} else if metaRes.JSON200 == nil || metaRes.JSON200.Jobs == nil || len(*metaRes.JSON200.Jobs) == 0 {
+		slog.Warn("job metadata response empty, falling back to default log paths")
+	} else {
+		meta := (*metaRes.JSON200.Jobs)[0]
+		if meta.StandardOutput != nil && *meta.StandardOutput != "" {
+			p := *meta.StandardOutput
+			if !path.IsAbs(p) {
+				p = path.Join(sessionPath, p)
+			}
+			c.stdoutPath = p
+		}
+		if meta.StandardError != nil && *meta.StandardError != "" {
+			p := *meta.StandardError
+			if !path.IsAbs(p) {
+				p = path.Join(sessionPath, p)
+			}
+			c.stderrPath = p
+		} else {
+			// Slurm writes stderr to stdout when StandardError is not explicitly set.
+			c.stderrPath = c.stdoutPath
+		}
+	}
+
+	// Save the state for recovery
+	if err := c.saveState(); err != nil {
 		return err
 	}
 
@@ -561,6 +601,14 @@ func (c *FirecrestRemoteSessionController) periodicSessionStatus(ctx context.Con
 				} else {
 					slog.Error("getCurrentStatus failed", "status", state, "error", err)
 				}
+				// Fetch any new session logs from the cluster filesystem.
+				c.fetchSessionLogs(childCtx)
+				// Persist offsets so we can resume after a restart.
+				if c.jobID != "" {
+					if err := c.saveState(); err != nil {
+						slog.Warn("failed to save controller state", "error", err)
+					}
+				}
 			}()
 		}
 	}
@@ -599,6 +647,88 @@ func (c *FirecrestRemoteSessionController) getCurrentStatus(ctx context.Context)
 	return state, nil
 }
 
+// fetchSessionLogs retrieves any new lines from the remote Slurm stdout/stderr files
+// via FirecREST and writes them to this container's stdout so they appear in kubectl logs.
+func (c *FirecrestRemoteSessionController) fetchSessionLogs(ctx context.Context) {
+	if c.jobID == "" || c.stdoutPath == "" {
+		return
+	}
+
+	for _, stream := range c.streamsToFetch() {
+		c.fetchLogStream(ctx, stream)
+	}
+}
+
+func (c *FirecrestRemoteSessionController) streamsToFetch() []string {
+	streams := []string{"stdout"}
+	if c.stderrPath != "" && c.stderrPath != c.stdoutPath {
+		streams = append(streams, "stderr")
+	}
+	return streams
+}
+
+func (c *FirecrestRemoteSessionController) fetchLogStream(ctx context.Context, stream string) {
+	var filePath string
+	var offset *int
+	var buf *bytes.Buffer
+
+	switch stream {
+	case "stdout":
+		filePath = c.stdoutPath
+		offset = &c.stdoutOffset
+		buf = &c.stdoutBuf
+	case "stderr":
+		filePath = c.stderrPath
+		offset = &c.stderrOffset
+		buf = &c.stderrBuf
+	default:
+		slog.Warn("unknown stream type", "stream", stream)
+		return
+	}
+
+	size := 524288 // 512KB
+	apiOffset := *offset + buf.Len()
+	params := GetViewFilesystemSystemNameOpsViewGetParams{
+		Path:   filePath,
+		Offset: &apiOffset,
+		Size:   &size,
+	}
+
+	res, err := c.client.GetViewFilesystemSystemNameOpsViewGetWithResponse(ctx, c.systemName, &params)
+	if err != nil {
+		slog.Warn("failed to fetch session log", "stream", stream, "error", err)
+		return
+	}
+	if res.JSON200 == nil || res.JSON200.Output == nil {
+		// File not available yet or empty — this is expected early in the job lifecycle.
+		return
+	}
+
+	content := *res.JSON200.Output
+	if content == "" {
+		return
+	}
+
+	// Append new content and flush any complete lines.
+	if _, err := buf.WriteString(content); err != nil {
+		slog.Warn("failed to buffer session log content", "stream", stream, "error", err)
+		return
+	}
+	for {
+		data := buf.Bytes()
+		idx := bytes.IndexByte(data, '\n')
+		if idx == -1 {
+			break
+		}
+		line := string(data[:idx])
+		if _, err := fmt.Fprintf(os.Stdout, "[session/%s] %s\n", stream, line); err != nil {
+			slog.Warn("failed to write session log", "stream", stream, "error", err)
+		}
+		buf.Next(idx + 1)
+		*offset += idx + 1
+	}
+}
+
 func (c *FirecrestRemoteSessionController) renderSessionScript(sessionScript string, fileSystems *[]FileSystem, secretsPath string) string {
 	return renderSessionScriptStatic(sessionScript, c.partition, fileSystems, secretsPath)
 }
@@ -673,7 +803,7 @@ func addSessionMountsToScript(sessionScript string, fileSystems *[]FileSystem, s
 	return strings.Replace(sessionScript, "#{{SESSION_MOUNTS_PLACEHOLDER}}", mountsStr, 1)
 }
 
-func (c *FirecrestRemoteSessionController) saveJobID() error {
+func (c *FirecrestRemoteSessionController) saveState() error {
 	if c.jobID == "" {
 		return fmt.Errorf("cannot save, job ID is not defined")
 	}
@@ -684,7 +814,11 @@ func (c *FirecrestRemoteSessionController) saveJobID() error {
 	savePath := c.getSavePath()
 
 	state := savedState{
-		JobID: c.jobID,
+		JobID:        c.jobID,
+		StdoutPath:   c.stdoutPath,
+		StderrPath:   c.stderrPath,
+		StdoutOffset: c.stdoutOffset,
+		StderrOffset: c.stderrOffset,
 	}
 	contents, err := json.Marshal(state)
 	if err != nil {
@@ -699,7 +833,7 @@ func (c *FirecrestRemoteSessionController) deleteSavedState() error {
 	return os.Remove(savePath)
 }
 
-func (c *FirecrestRemoteSessionController) recoverJobID() error {
+func (c *FirecrestRemoteSessionController) recoverState() error {
 	contents, err := os.ReadFile(c.getSavePath())
 	if err != nil {
 		return nil
@@ -714,11 +848,19 @@ func (c *FirecrestRemoteSessionController) recoverJobID() error {
 		c.jobID = state.JobID
 		slog.Info("recovered job ID", "jobID", c.jobID)
 	}
+	c.stdoutPath = state.StdoutPath
+	c.stderrPath = state.StderrPath
+	c.stdoutOffset = state.StdoutOffset
+	c.stderrOffset = state.StderrOffset
 	return nil
 }
 
 type savedState struct {
-	JobID string `json:"job_id"`
+	JobID        string `json:"job_id"`
+	StdoutPath   string `json:"stdout_path,omitempty"`
+	StderrPath   string `json:"stderr_path,omitempty"`
+	StdoutOffset int    `json:"stdout_offset,omitempty"`
+	StderrOffset int    `json:"stderr_offset,omitempty"`
 }
 
 func (c *FirecrestRemoteSessionController) getSaveDirPath() string {

internal/remote/firecrest/controller_test.go

@@ -59,3 +59,26 @@ func TestRenderSessionScriptStatic(t *testing.T) {
 	assert.Contains(t, foundMounts, "\"/users:/home/users:ro\"")
 	assert.Contains(t, foundMounts, "\"/secrets:/secrets:ro\"")
 }
+
+func TestStreamsToFetch(t *testing.T) {
+	tests := []struct {
+		name       string
+		stdoutPath string
+		stderrPath string
+		want       []string
+	}{
+		{"different paths", "/out", "/err", []string{"stdout", "stderr"}},
+		{"same path", "/out", "/out", []string{"stdout"}},
+		{"empty stderr", "/out", "", []string{"stdout"}},
+		{"stderr explicitly eq stdout", "/out", "/out", []string{"stdout"}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := &FirecrestRemoteSessionController{
+				stdoutPath: tt.stdoutPath,
+				stderrPath: tt.stderrPath,
+			}
+			assert.Equal(t, tt.want, c.streamsToFetch())
+		})
+	}
+}