fix(dataset push): make live ingestion work end-to-end

internal/cli/dataset.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"path/filepath"
 
 	"github.com/spf13/cobra"
 	"gopkg.in/yaml.v3"
@@ -74,6 +75,7 @@ func newDatasetPushCmd() *cobra.Command {
  category string
  intent string
  labelColumn string
+ targetSize string
 
  // Operations flags.
  dryRun bool
@@ -145,6 +147,7 @@ Exit codes:
  Context: contextOverride,
  Namespace: nsOverride,
  Spec: push.SpecArgs{Table: table, Category: category, Intent: intent, LabelColumn: labelColumn},
+ TargetSizeFlag: targetSize,
  DryRun: dryRun,
  IngestorSAName: ingestorSAName,
  StagePodImage: stagePodImage,
@@ -175,6 +178,9 @@ Exit codes:
  "intent: train|test")
 	cmd.Flags().StringVar(&labelColumn, "label-column", "",
  "column name in labels.csv that holds the label")
+	cmd.Flags().StringVar(&targetSize, "target-size", "",
+ "image resolution as WxH (e.g. 512x512). Default: auto-detected from the first image. "+
+ "All images must share this resolution — the ingestor validates it, it does not resize.")
 
 	cmd.Flags().BoolVar(&dryRun, "dry-run", false,
  "validate + discover + walk, but don't create any cluster resources")
@@ -209,6 +215,7 @@ type runDatasetPushArgs struct {
 	Context string
 	Namespace string
 	Spec push.SpecArgs
+	TargetSizeFlag string // raw --target-size; resolved after Discover
 	DryRun bool
 	IngestorSAName string
 	StagePodImage string
@@ -258,11 +265,48 @@ func runDatasetPush(ctx context.Context, out, errOut io.Writer, a runDatasetPush
  "tracebloc/client#147 non-goals.", a.Spec.Category)}
 	}
 
-	// 3. Synthesize the spec from flags + validate against schema.
+	// 3. Walk the local directory FIRST. Enforces layout + size caps,
+	// and gives us the image list the target-size auto-detect below
+	// needs. Both this and the schema check are local "fail fast"
+	// steps; doing the walk first lets the synthesized spec carry
+	// the resolved target_size.
+	layout, err := push.Discover(a.LocalPath)
+	if err != nil {
+ return &exitError{code: 3, err: err}
+	}
+
+	// 3a. Resolve the image target resolution. The ingestor's
+	// image_classification default is 512x512 and it VALIDATES
+	// (it does not resize), so a mismatch hard-fails the run with
+	// an "incorrect resolution" error. Honour an explicit
+	// --target-size; otherwise auto-detect from the first image so
+	// the common "all my images are NxN" case just works without
+	// the customer needing to know the knob exists.
+	if a.TargetSizeFlag != "" {
+ w, h, perr := push.ParseTargetSize(a.TargetSizeFlag)
+ if perr != nil {
+ return &exitError{code: 2, err: perr}
+ }
+ a.Spec.TargetSize = []int{w, h}
+	} else if len(layout.Images) > 0 {
+ if w, h, derr := push.DetectImageSize(layout.Images[0]); derr == nil {
+ a.Spec.TargetSize = []int{w, h}
+ _, _ = fmt.Fprintf(out,
+ "Auto-detected image target size %dx%d from %s (override with --target-size).\n",
+ w, h, filepath.Base(layout.Images[0]))
+ } else {
+ _, _ = fmt.Fprintf(errOut,
+ "Note: couldn't auto-detect image size (%v); using the ingestor "+
+ "default. Pass --target-size WxH if ingestion reports a "+
+ "resolution mismatch.\n", derr)
+ }
+	}
+
+	// 4. Synthesize the spec from flags + validate against schema.
 	// Catches "bad category", "missing intent" etc. BEFORE we
-	// touch the filesystem or the cluster. The error formatter
-	// is the same one ingest validate uses, so a customer who
-	// YAML'd manually first sees identical wording.
+	// touch the cluster. The error formatter is the same one
+	// ingest validate uses, so a customer who YAML'd manually
+	// first sees identical wording.
 	spec := a.Spec.Build()
 	specBytes, err := yaml.Marshal(spec)
 	if err != nil {
@@ -295,14 +339,6 @@ func runDatasetPush(ctx context.Context, out, errOut io.Writer, a runDatasetPush
  return &exitError{code: 2, err: errors.New("synthesized spec failed schema validation; check the flag values above")}
 	}
 
-	// 4. Walk the local directory. Enforces layout + size caps;
-	// customer sees a clear pointer to expected layout if they
-	// pass the wrong directory.
-	layout, err := push.Discover(a.LocalPath)
-	if err != nil {
- return &exitError{code: 3, err: err}
-	}
-
 	// 5. Cluster discovery — same kubeconfig path as `cluster info`.
 	// Errors mirror that command's exit-code contract (3 for
 	// kubeconfig, 4 for missing release) so behaviour is
@@ -527,13 +563,12 @@ func printPushPreflight(
 	_, _ = fmt.Fprintf(out, " category: %s\n", spec["category"])
 	_, _ = fmt.Fprintf(out, " intent: %s\n", spec["intent"])
 	_, _ = fmt.Fprintf(out, " label column: %s\n", spec["label"])
-	_, _ = fmt.Fprintf(out, " destination: %s\n", push.StagedPrefix(spec["table"].(string)))
+	_, _ = fmt.Fprintf(out, " destination: %s\n", push.FinalDestPrefix(spec["table"].(string)))
 	_, _ = fmt.Fprintln(out)
 
 	if !dryRun {
- _, _ = fmt.Fprintf(out, "Next: stage %d files (%s) → %s\n",
- 1+len(layout.Images), push.HumanBytes(layout.TotalBytes),
- push.StagedPrefix(spec["table"].(string)))
+ _, _ = fmt.Fprintf(out, "Next: stage %d files (%s) for table %q\n",
+ 1+len(layout.Images), push.HumanBytes(layout.TotalBytes), spec["table"])
  _, _ = fmt.Fprintln(out)
 	}
 }

internal/push/detect.go

@@ -0,0 +1,70 @@
+package push
+
+import (
+	"fmt"
+	"image"
+	"os"
+	"strconv"
+	"strings"
+
+	// Register the stdlib image decoders so image.DecodeConfig can
+	// read the headers of the formats the image_classification layout
+	// accepts. webp is NOT in the stdlib — DetectImageSize returns an
+	// error for it and the caller falls back to requiring
+	// --target-size. (.jpg/.jpeg both decode via image/jpeg.)
+	_ "image/gif"
+	_ "image/jpeg"
+	_ "image/png"
+)
+
+// DetectImageSize returns the pixel width and height of the image at
+// path by decoding only its header (image.DecodeConfig — it does not
+// read the pixel data, so it's cheap even for large images).
+//
+// Supports the stdlib-registered formats (jpeg, png, gif). Returns an
+// error for formats without a registered decoder (notably webp); the
+// caller treats that as "couldn't auto-detect" and falls back to the
+// ingestor default, advising --target-size.
+func DetectImageSize(path string) (width, height int, err error) {
+	f, err := os.Open(path)
+	if err != nil {
+ return 0, 0, err
+	}
+	defer func() { _ = f.Close() }()
+
+	cfg, _, err := image.DecodeConfig(f)
+	if err != nil {
+ return 0, 0, fmt.Errorf("decoding image header %q: %w", path, err)
+	}
+	return cfg.Width, cfg.Height, nil
+}
+
+// ParseTargetSize parses a --target-size flag value into [width,
+// height]. Accepts "WxH" (the documented form, e.g. "512x512") and
+// "W,H" as a convenience. Both dimensions must be positive integers.
+func ParseTargetSize(s string) (width, height int, err error) {
+	sep := "x"
+	if strings.Contains(s, ",") {
+ sep = ","
+	}
+	parts := strings.Split(s, sep)
+	if len(parts) != 2 {
+ return 0, 0, fmt.Errorf(
+ "target size %q must be WxH (e.g. 512x512)", s)
+	}
+	width, err = strconv.Atoi(strings.TrimSpace(parts[0]))
+	if err != nil {
+ return 0, 0, fmt.Errorf(
+ "target size %q: width is not an integer: %w", s, err)
+	}
+	height, err = strconv.Atoi(strings.TrimSpace(parts[1]))
+	if err != nil {
+ return 0, 0, fmt.Errorf(
+ "target size %q: height is not an integer: %w", s, err)
+	}
+	if width <= 0 || height <= 0 {
+ return 0, 0, fmt.Errorf(
+ "target size %q: width and height must both be positive", s)
+	}
+	return width, height, nil
+}

internal/push/detect_test.go

@@ -0,0 +1,88 @@
+package push
+
+import (
+	"image"
+	"image/png"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+// TestParseTargetSize covers the --target-size flag parser: the
+// documented WxH form, the W,H convenience form, and the rejection
+// cases (missing dimension, non-integer, non-positive, wrong arity).
+func TestParseTargetSize(t *testing.T) {
+	cases := []struct {
+ in string
+ w, h int
+ wantErr bool
+	}{
+ {"512x512", 512, 512, false},
+ {"640x480", 640, 480, false},
+ {"512,512", 512, 512, false},
+ {"1x1", 1, 1, false},
+ {"512", 0, 0, true},
+ {"512x", 0, 0, true},
+ {"x512", 0, 0, true},
+ {"0x512", 0, 0, true},
+ {"-4x512", 0, 0, true},
+ {"512x512x512", 0, 0, true},
+ {"abcxdef", 0, 0, true},
+ {"", 0, 0, true},
+	}
+	for _, c := range cases {
+ w, h, err := ParseTargetSize(c.in)
+ if c.wantErr {
+ if err == nil {
+ t.Errorf("ParseTargetSize(%q) = (%d,%d,nil), want error", c.in, w, h)
+ }
+ continue
+ }
+ if err != nil {
+ t.Errorf("ParseTargetSize(%q) unexpected error: %v", c.in, err)
+ continue
+ }
+ if w != c.w || h != c.h {
+ t.Errorf("ParseTargetSize(%q) = (%d,%d), want (%d,%d)", c.in, w, h, c.w, c.h)
+ }
+	}
+}
+
+// TestDetectImageSize_PNG: a real (generated) PNG's header is decoded
+// to its true dimensions. Pins the auto-detect path used when the
+// customer doesn't pass --target-size.
+func TestDetectImageSize_PNG(t *testing.T) {
+	dir := t.TempDir()
+	p := filepath.Join(dir, "img.png")
+	f, err := os.Create(p)
+	if err != nil {
+ t.Fatal(err)
+	}
+	if err := png.Encode(f, image.NewRGBA(image.Rect(0, 0, 320, 200))); err != nil {
+ t.Fatal(err)
+	}
+	_ = f.Close()
+
+	w, h, err := DetectImageSize(p)
+	if err != nil {
+ t.Fatalf("DetectImageSize: %v", err)
+	}
+	if w != 320 || h != 200 {
+ t.Errorf("DetectImageSize = (%d,%d), want (320,200)", w, h)
+	}
+}
+
+// TestDetectImageSize_Unsupported: a non-image (or unregistered
+// format) returns an error so the caller falls back to the ingestor
+// default + advises --target-size, rather than silently using a
+// bogus size.
+func TestDetectImageSize_Unsupported(t *testing.T) {
+	dir := t.TempDir()
+	p := filepath.Join(dir, "note.txt")
+	if err := os.WriteFile(p, []byte("not an image"), 0o644); err != nil {
+ t.Fatal(err)
+	}
+	if _, _, err := DetectImageSize(p); err == nil {
+ t.Error("DetectImageSize on non-image returned nil error; want a decode error")
+	}
+}