Improve mass output downloads (#135)

* AllRuns takes precedence over any specified number of runs

* Return unmatched identifiers while filtering sub jobs instead of returning an error

* Get the complete run details while downloading outputs if a simplified run object that doesn't contain the workflow version ID was used initially

* Do not error out on the first failed run and print extra details in the download status messages

* Return the destination path while downloading outputs and pass the individual download error details to the caller

Author: mhmdiaa

cmd/execute/execute.go

@@ -295,11 +295,11 @@ func run(cfg *Config) error {
 		if err != nil {
 			return fmt.Errorf("failed to get run: %w", err)
 		}
-		results, err := actions.DownloadRunOutput(client, run, cfg.NodesToDownload, []string{}, cfg.OutputDirectory)
+		results, runDir, err := actions.DownloadRunOutput(client, run, cfg.NodesToDownload, []string{}, cfg.OutputDirectory)
 		if err != nil {
 			return fmt.Errorf("failed to download run outputs: %w", err)
 		}
-		actions.PrintDownloadResults(results)
+		actions.PrintDownloadResults(results, *run.ID, runDir)
 	}
 
 	return nil

cmd/output/output.go

@@ -83,12 +83,24 @@ func run(cfg *Config) error {
 	files := cfg.GetFiles()
 	path := cfg.GetOutputPath()
 
-	for _, run := range runs {
-		results, err := actions.DownloadRunOutput(client, &run, nodes, files, path)
+	var errorCount int
+	for i, run := range runs {
+		fmt.Fprintf(os.Stderr, "Downloading outputs for run %s\n", run.ID.String())
+		results, runDir, err := actions.DownloadRunOutput(client, &run, nodes, files, path)
 		if err != nil {
-			return fmt.Errorf("failed to download run output: %w", err)
+			fmt.Fprintf(os.Stderr, "Warning: failed to download run output for run %s: %v\n", run.ID.String(), err)
+			errorCount++
 		}
-		actions.PrintDownloadResults(results)
+
+		actions.PrintDownloadResults(results, *run.ID, runDir)
+		// Print a newline visual separation if there are more runs to process
+		if i < len(runs)-1 {
+			fmt.Fprintln(os.Stderr)
+		}
+	}
+
+	if errorCount == len(runs) {
+		return fmt.Errorf("failed to download outputs for all runs")
 	}
 	return nil
 }

cmd/stop/stop.go

@@ -124,9 +124,12 @@ func run(cfg *Config) error {
 		}
 		subJobs = trickest.LabelSubJobs(subJobs, *version)
 
-		matchingSubJobs, err := trickest.FilterSubJobs(subJobs, cfg.Nodes)
-		if err != nil {
-			return fmt.Errorf("no running nodes matching your query were found in the run %s: %w", run.ID.String(), err)
+		matchingSubJobs, unmatchedNodes := trickest.FilterSubJobs(subJobs, cfg.Nodes)
+		if len(matchingSubJobs) == 0 {
+			return fmt.Errorf("no running nodes matching your query %q were found in the run %s", strings.Join(cfg.Nodes, ","), run.ID.String())
+		}
+		if len(unmatchedNodes) > 0 {
+			fmt.Fprintf(os.Stderr, "Warning: The following nodes were not found in run %s: %s. Proceeding with the remaining nodes\n", run.ID.String(), strings.Join(unmatchedNodes, ","))
 		}
 
 		for _, subJob := range matchingSubJobs {

pkg/actions/output.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 	"sync"
 
 	"github.com/google/uuid"
@@ -19,7 +20,7 @@ type DownloadResult struct {
 	Error      error
 }
 
-func PrintDownloadResults(results []DownloadResult) {
+func PrintDownloadResults(results []DownloadResult, runID uuid.UUID, destinationPath string) {
 	successCount := 0
 	failureCount := 0
 	for _, result := range results {
@@ -28,46 +29,58 @@ func PrintDownloadResults(results []DownloadResult) {
 		} else {
 			failureCount++
 			if result.FileName != "" {
-				fmt.Fprintf(os.Stderr, "Warning: Failed to download file %q for node %q: %v\n", result.FileName, result.SubJobName, result.Error)
+				fmt.Fprintf(os.Stderr, "Warning: Failed to download file %q for node %q in run %s: %v\n", result.FileName, result.SubJobName, runID.String(), result.Error)
 			} else {
-				fmt.Fprintf(os.Stderr, "Warning: Failed to download output for node %q: %v\n", result.SubJobName, result.Error)
+				fmt.Fprintf(os.Stderr, "Warning: Failed to download output for node %q in run %s: %v\n", result.SubJobName, runID.String(), result.Error)
 			}
 		}
 	}
 
 	if failureCount > 0 {
-		fmt.Fprintf(os.Stderr, "Download completed with %d successful and %d failed downloads\n", successCount, failureCount)
+		fmt.Fprintf(os.Stderr, "Download completed with %d successful and %d failed downloads for run %s into %q\n", successCount, failureCount, runID.String(), destinationPath+"/")
 	} else if successCount > 0 {
-		fmt.Printf("Successfully downloaded outputs for %d sub-jobs\n", successCount)
+		fmt.Printf("Successfully downloaded %d outputs from run %s into %q\n", successCount, runID.String(), destinationPath+"/")
 	}
 }
 
-func DownloadRunOutput(client *trickest.Client, run *trickest.Run, nodes []string, files []string, destinationPath string) ([]DownloadResult, error) {
+// DownloadRunOutput downloads the outputs for the specified nodes in the run
+// Returns the download result summary, the directory where the outputs were saved, and an error if _all_ of the downloads failed
+func DownloadRunOutput(client *trickest.Client, run *trickest.Run, nodes []string, files []string, destinationPath string) ([]DownloadResult, string, error) {
 	if run.Status == "PENDING" || run.Status == "SUBMITTED" {
-		return nil, fmt.Errorf("run %s has not started yet (status: %s)", run.ID.String(), run.Status)
+		return nil, "", fmt.Errorf("run %s has not started yet (status: %s)", run.ID.String(), run.Status)
 	}
 
 	ctx := context.Background()
 
 	subJobs, err := client.GetSubJobs(ctx, *run.ID)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get subjobs for run %s: %w", run.ID.String(), err)
+		return nil, "", fmt.Errorf("failed to get subjobs for run %s: %w", run.ID.String(), err)
 	}
 
+	// If the run was retrieved through the GetRuns() method, the WorkflowVersionInfo field will be nil
+	if run.WorkflowVersionInfo == nil {
+		run, err = client.GetRun(ctx, *run.ID)
+		if err != nil {
+			return nil, "", fmt.Errorf("failed to get run details for run %s: %w", run.ID.String(), err)
+		}
+	}
 	version, err := client.GetWorkflowVersion(ctx, *run.WorkflowVersionInfo)
 	if err != nil {
-		return nil, fmt.Errorf("could not get workflow version for run %s: %w", run.ID.String(), err)
+		return nil, "", fmt.Errorf("could not get workflow version for run %s: %w", run.ID.String(), err)
 	}
 	subJobs = trickest.LabelSubJobs(subJobs, *version)
 
-	matchingSubJobs, err := trickest.FilterSubJobs(subJobs, nodes)
-	if err != nil {
-		return nil, fmt.Errorf("no completed node outputs matching your query were found in the run %s: %w", run.ID.String(), err)
+	matchingSubJobs, unmatchedNodes := trickest.FilterSubJobs(subJobs, nodes)
+	if len(matchingSubJobs) == 0 {
+		return nil, "", fmt.Errorf("no completed node outputs matching your query %q were found in the run %s", strings.Join(nodes, ","), run.ID.String())
+	}
+	if len(unmatchedNodes) > 0 {
+		fmt.Fprintf(os.Stderr, "Warning: The following nodes were not found in run %s: %s. Proceeding with the remaining nodes\n", run.ID.String(), strings.Join(unmatchedNodes, ","))
 	}
 
 	runDir, err := filesystem.CreateRunDir(destinationPath, *run)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create directory for run %s: %w", run.ID.String(), err)
+		return nil, "", fmt.Errorf("failed to create directory for run %s: %w", run.ID.String(), err)
 	}
 
 	var allResults []DownloadResult
@@ -84,13 +97,7 @@ func DownloadRunOutput(client *trickest.Client, run *trickest.Run, nodes []strin
 		allResults = append(allResults, results...)
 	}
 
-	// If all subjobs failed, return an error
-	if errCount == len(allResults) && len(allResults) > 0 {
-		return allResults, fmt.Errorf("failed to download outputs for all nodes")
-	}
-
-	// If only some failed, return results but no error
-	return allResults, nil
+	return allResults, runDir, nil
 }
 
 func downloadSubJobOutput(client *trickest.Client, savePath string, subJob *trickest.SubJob, files []string, runID *uuid.UUID, isModule bool) []DownloadResult {

pkg/config/workflowrunspec.go

@@ -49,7 +49,7 @@ func (s WorkflowRunSpec) GetRuns(ctx context.Context, client *trickest.Client) (
 	}
 
 	limit := 0 // 0 means get all runs
-	if s.NumberOfRuns > 0 {
+	if s.NumberOfRuns > 0 && !s.AllRuns {
 		limit = s.NumberOfRuns
 	}
 

pkg/trickest/subjob.go

@@ -172,35 +172,36 @@ func LabelSubJobs(subJobs []SubJob, version WorkflowVersion) []SubJob {
 }
 
 // FilterSubJobs filters the subjobs based on the identifiers (label or name (node ID))
-func FilterSubJobs(subJobs []SubJob, identifiers []string) ([]SubJob, error) {
+// Returns matched subjobs and unmatched identifiers
+func FilterSubJobs(subJobs []SubJob, identifiers []string) ([]SubJob, []string) {
 	if len(identifiers) == 0 {
 		return subJobs, nil
 	}
 
-	var foundNodes []string
 	var matchingSubJobs []SubJob
+	matchedIdentifiers := make(map[string]bool)
 
 	for _, subJob := range subJobs {
 		labelExists := slices.Contains(identifiers, subJob.Label)
 		nameExists := slices.Contains(identifiers, subJob.Name)
 
-		if labelExists {
-			foundNodes = append(foundNodes, subJob.Label)
-		}
-		if nameExists {
-			foundNodes = append(foundNodes, subJob.Name)
-		}
-
 		if labelExists || nameExists {
 			matchingSubJobs = append(matchingSubJobs, subJob)
+			if labelExists {
+				matchedIdentifiers[subJob.Label] = true
+			}
+			if nameExists {
+				matchedIdentifiers[subJob.Name] = true
+			}
 		}
 	}
 
+	var unmatchedIdentifiers []string
 	for _, identifier := range identifiers {
-		if !slices.Contains(foundNodes, identifier) {
-			return nil, fmt.Errorf("subjob with name or label %s not found", identifier)
+		if !matchedIdentifiers[identifier] {
+			unmatchedIdentifiers = append(unmatchedIdentifiers, identifier)
 		}
 	}
 
-	return matchingSubJobs, nil
+	return matchingSubJobs, unmatchedIdentifiers
 }