cli: refactor get-batch CLI

alex-aizman · alex-aizman · commit 2f18344e3c55 · 2025-07-18T11:31:56.000-04:00
* split request parsing in two parts:
  - parse/validate CLI flags and arguments
  - populate `apc.MossReq` from native spec or Lhotse manifest
* output file:
  - clarify and maybe warn when Lhotse multi-batch
  - enforce for single-batch
* with refactoring

Signed-off-by: Alex Aizman &lt;alex.aizman@gmail.com&gt;
diff --git a/cmd/cli/cli/lhotse.go b/cmd/cli/cli/lhotse.go
@@ -230,7 +230,7 @@ func parseLhotseSrc(uri string) (bck cmn.Bck, objName string, err error) {
 // generate multiple batches in a streaming fashion
 //
 
-func lhotseMultiBatch(c *cli.Context, outCtx *mossReqParseCtx) error {
+func lhotseMultiBatch(c *cli.Context, outCtx *mossReqCtx) error {
 	inCtx, err := openLhotseReader(c)
 	if err != nil {
 		return err
@@ -310,7 +310,7 @@ func lhotseMultiBatch(c *cli.Context, outCtx *mossReqParseCtx) error {
 	return nil
 }
 
-func genLhotseBatch(c *cli.Context, outCtx *mossReqParseCtx, batch []apc.MossIn, outputFile string) error {
+func genLhotseBatch(c *cli.Context, outCtx *mossReqCtx, batch []apc.MossIn, outputFile string) error {
 	// 1) create MossReq for this batch
 	req := apc.MossReq{
 		In:            batch,
diff --git a/cmd/cli/cli/ml.go b/cmd/cli/cli/ml.go
@@ -91,22 +91,78 @@ var (
 	}
 )
 
-type (
-	mossReqParseCtx struct {
-		req       apc.MossReq
-		outFile   string
-		bck       cmn.Bck
-		pt        *cos.ParsedTemplate // valid when batchSize > 0
-		batchSize int
+type mossReqCtx struct {
+	// get-batch API
+	req apc.MossReq
+	bck cmn.Bck
+
+	// output file
+	outFile string
+
+	// additional/optional command line input
+	names []string
+	shift int
+
+	// Lhotse-specific
+	pt        *cos.ParsedTemplate
+	batchSize int
+}
+
+func getBatchHandler(c *cli.Context) error {
+	ctx, err := parseMlArgs(c)
+	if err != nil {
+		return err
+	}
+	if err := ctx.load(c); err != nil {
+		return err
+	}
+
+	if ctx.batchSize > 0 {
+		debug.Assert(flagIsSet(c, lhotseManifestFlag), "native (non-lhotse) batching not implemented yet")
+		return lhotseMultiBatch(c, ctx)
+	}
+	outFile, bck := ctx.outFile, ctx.bck
+
+	// output
+	w, wfh, err := createDstFile(c, outFile, false /*allow stdout*/)
+	if err != nil {
+		if err == errUserCancel {
+			return nil
+		}
+		return err
+	}
+
+	// do
+	req := &ctx.req
+	resp, err := api.GetBatch(apiBP, bck, req, w)
+
+	if wfh != nil {
+		wfh.Close()
+	}
+	if err == nil {
+		debug.Assert(req.StreamingGet || len(resp.Out) == len(req.In))
+
+		if !flagIsSet(c, nonverboseFlag) {
+			var msg string
+			if req.StreamingGet {
+				msg = fmt.Sprintf("Streamed %d objects to %s", len(req.In), outFile)
+			} else {
+				msg = fmt.Sprintf("Created ML batch archive %s with %d objects", outFile, len(resp.Out))
+			}
+			actionDone(c, msg)
+		}
+		return nil
 	}
-)
 
-// TODO -- FIXME: refactor - split in parts
-func buildMossReq(c *cli.Context) (*mossReqParseCtx, error) {
-	var (
-		req   apc.MossReq
-		shift int
-	)
+	// cleanup
+	if wfh != nil {
+		cos.RemoveFile(outFile)
+	}
+	return err
+}
+
+// parse CLI arguments and setup context
+func parseMlArgs(c *cli.Context) (ctx *mossReqCtx, _ error) {
 	if err := errMutuallyExclusive(c, listFlag, templateFlag, specFlag); err != nil {
 		return nil, err
 	}
@@ -118,9 +174,6 @@ func buildMossReq(c *cli.Context) (*mossReqParseCtx, error) {
 	}
 
 	// bucket [+list|template]
-	var (
-		names []string
-	)
 	bck, objNameOrTmpl, err := parseBckObjURI(c, c.Args().Get(0), true /*emptyObjnameOK*/)
 	if flagIsSet(c, listFlag) && err != nil {
 		return nil, fmt.Errorf("option %s requires bucket in the command line [err: %v]", qflprn(listFlag), err)
@@ -129,8 +182,9 @@ func buildMossReq(c *cli.Context) (*mossReqParseCtx, error) {
 		return nil, fmt.Errorf("option %s requires bucket in the command line [err: %v]", qflprn(templateFlag), err)
 	}
 
+	ctx = &mossReqCtx{bck: bck}
 	if err == nil {
-		shift++
+		ctx.shift++
 		oltp, err := dopOLTP(c, bck, objNameOrTmpl)
 		if err != nil {
 			return nil, err
@@ -144,99 +198,104 @@ func buildMossReq(c *cli.Context) (*mossReqParseCtx, error) {
 			if err := pt.CheckIsRange(); err != nil {
 				return nil, err
 			}
-			if names, err = pt.Expand(); err != nil {
+			if ctx.names, err = pt.Expand(); err != nil {
 				return nil, err
 			}
 		case oltp.list == "":
 			if objNameOrTmpl != "" {
-				names = []string{objNameOrTmpl}
+				ctx.names = []string{objNameOrTmpl}
 			}
 		default:
-			names = splitCsv(oltp.list)
+			ctx.names = splitCsv(oltp.list)
 		}
 	}
 
-	if len(names) == 0 && !flagIsSet(c, specFlag) && !flagIsSet(c, lhotseManifestFlag) {
+	if len(ctx.names) == 0 && !flagIsSet(c, specFlag) && !flagIsSet(c, lhotseManifestFlag) {
 		return nil, fmt.Errorf("with no (%s, %s) options expecting object names and/or archived filenames in the command line",
 			qflprn(specFlag), qflprn(lhotseManifestFlag))
 	}
 
-	// native spec
-	if flagIsSet(c, specFlag) {
-		specBytes, ext, err := loadSpec(c)
+	// lhotse only: given (batchSizeFlag & outputTemplateFlag) outFile(s) are computed from the latter
+	if c.NArg() > ctx.shift {
+		// Note: may be unused in Lhotse multi-batch mode
+		ctx.outFile = c.Args().Get(ctx.shift)
+		outputFormat, err := archive.Strict("", cos.Ext(ctx.outFile))
 		if err != nil {
 			return nil, err
 		}
-		if err := parseSpec(ext, specBytes, &req); err != nil {
-			return nil, err
-		}
+		ctx.req.OutputFormat = outputFormat
 	}
 
-	var (
-		outputFormat string
-		outFile      string
-	)
-	// given (batchSizeFlag and outputTemplateFlag)
-	// outFile can be computed from the latter
-	// (currently, lhotse only)
-	if c.NArg() > shift {
-		outFile = c.Args().Get(shift)
-		outputFormat, err = archive.Strict("", cos.Ext(outFile))
+	// there's no real way to check these assorted overrides; common expectation, though,
+	// is for command line to take precedence
+
+	ctx.req.ContinueOnErr = flagIsSet(c, continueOnErrorFlag)
+	ctx.req.StreamingGet = flagIsSet(c, streamingGetFlag)
+	ctx.req.OnlyObjName = flagIsSet(c, omitSrcBucketNameFlag)
+
+	return ctx, nil
+}
+
+// populate get-batch request from native (single-batch) spec or Lhotse manifest
+func (ctx *mossReqCtx) load(c *cli.Context) error {
+	// native spec
+	if flagIsSet(c, specFlag) {
+		specBytes, ext, err := loadSpec(c)
 		if err != nil {
-			return nil, err
+			return err
+		}
+
+		outputFormat := ctx.req.OutputFormat
+		if err := parseSpec(ext, specBytes, &ctx.req); err != nil {
+			return err
 		}
-		if req.OutputFormat != "" && outputFormat != "" && req.OutputFormat != outputFormat {
+		// warn
+		if ctx.req.OutputFormat != "" && outputFormat != "" && ctx.req.OutputFormat != outputFormat {
 			if !flagIsSet(c, nonverboseFlag) {
 				warn := fmt.Sprintf("output format %s in the command line takes precedence (over %s specified %s)",
-					outputFormat, qflprn(specFlag), req.OutputFormat)
+					outputFormat, qflprn(specFlag), ctx.req.OutputFormat)
 				actionWarn(c, warn)
 			}
 		}
-		req.OutputFormat = outputFormat
 	}
 
-	// NOTE: no real way to check these assorted overrides; common expectation, though,
-	// is for command line to take precedence
-	req.ContinueOnErr = flagIsSet(c, continueOnErrorFlag)
-	req.StreamingGet = flagIsSet(c, streamingGetFlag)
-	req.OnlyObjName = flagIsSet(c, omitSrcBucketNameFlag)
-
-	var (
-		ctx = mossReqParseCtx{outFile: outFile, bck: bck}
-	)
 	// lhotse spec
 	if flagIsSet(c, lhotseManifestFlag) {
+		var err error
 		ctx.batchSize, ctx.pt, err = parseLhotseBatchFlags(c)
 		if err != nil {
-			return nil, err
+			return err
 		}
 		if ctx.batchSize > 0 {
 			// note: early return to generate multiple apc.MossReq requests and batches
-			return &ctx, nil
+			if ctx.outFile != "" {
+				warn := fmt.Sprintf("output file %s is ignored in multi-batch mode (batch size %d)", ctx.outFile, ctx.batchSize)
+				actionWarn(c, warn)
+			}
+			return nil
 		}
-
-		if outFile == "" {
-			return nil, missingArgumentsError(c, c.Command.ArgsUsage)
+		if ctx.outFile == "" {
+			return errors.New("output file is required for single-batch mode")
 		}
 
 		ins, err := loadAndParseLhotse(c)
 		if err != nil {
-			return nil, err
+			return err
 		}
-		req.In = ins
-	} else if outFile == "" {
-		return nil, missingArgumentsError(c, c.Command.ArgsUsage)
+		ctx.req.In = ins
+	} else if ctx.outFile == "" {
+		return missingArgumentsError(c, c.Command.ArgsUsage)
 	}
 
-	if len(names) > 0 {
-		if len(req.In) == 0 {
-			req.In = make([]apc.MossIn, 0, len(names))
+	if len(ctx.names) > 0 {
+		if len(ctx.req.In) == 0 {
+			ctx.req.In = make([]apc.MossIn, 0, len(ctx.names))
 		} else {
 			warn := fmt.Sprintf("adding %d command-line defined name%s to the %d spec-defined",
-				len(names), cos.Plural(len(names)), len(req.In))
+				len(ctx.names), cos.Plural(len(ctx.names)), len(ctx.req.In))
 			actionWarn(c, warn)
 		}
-		for _, o := range names {
+		for _, o := range ctx.names {
 			in := apc.MossIn{
 				ObjName: o,
 				// no need to insert command-line bck -
@@ -246,64 +305,13 @@ func buildMossReq(c *cli.Context) (*mossReqParseCtx, error) {
 				in.ObjName = oname
 				in.ArchPath = archpath
 			}
-			req.In = append(req.In, in)
-		}
-	}
-
-	if len(req.In) == 0 {
-		return nil, errors.New("empty get-batch request")
-	}
-
-	ctx.req = req
-	return &ctx, nil
-}
-
-func getBatchHandler(c *cli.Context) error {
-	ctx, err := buildMossReq(c)
-	if err != nil {
-		return err
-	}
-
-	if ctx.batchSize > 0 {
-		debug.Assert(flagIsSet(c, lhotseManifestFlag), "native (non-lhotse) batching not implemented yet")
-		return lhotseMultiBatch(c, ctx)
-	}
-	outFile, bck := ctx.outFile, ctx.bck
-
-	// output
-	w, wfh, err := createDstFile(c, outFile, false /*allow stdout*/)
-	if err != nil {
-		if err == errUserCancel {
-			return nil
+			ctx.req.In = append(ctx.req.In, in)
 		}
-		return err
-	}
-
-	// do
-	req := &ctx.req
-	resp, err := api.GetBatch(apiBP, bck, req, w)
-
-	if wfh != nil {
-		wfh.Close()
 	}
-	if err == nil {
-		debug.Assert(req.StreamingGet || len(resp.Out) == len(req.In))
 
-		if !flagIsSet(c, nonverboseFlag) {
-			var msg string
-			if req.StreamingGet {
-				msg = fmt.Sprintf("Streamed %d objects to %s", len(req.In), outFile)
-			} else {
-				msg = fmt.Sprintf("Created ML batch archive %s with %d objects", outFile, len(resp.Out))
-			}
-			actionDone(c, msg)
-		}
-		return nil
+	if len(ctx.req.In) == 0 {
+		return errors.New("empty get-batch request")
 	}
 
-	// cleanup
-	if wfh != nil {
-		cos.RemoveFile(outFile)
-	}
-	return err
+	return nil
 }
