fix(data ingest): reclaim the staged source copy after a clean success (#166)

internal/cli/data.go

@@ -828,6 +828,10 @@ other collaborators train against it without ever seeing the raw files.`))
 	// whose status matches the exit code. (Bugbot #38.)
 	status, exitErr := classifyPushOutcome(submitRes, err)
 
+	// Emit the machine-readable result BEFORE the best-effort staging
+	// reclaim below, so a scripted --output-json consumer gets its result
+	// object at ingest-completion latency and never waits on a slow
+	// cluster-side cleanup that has no bearing on the ingest outcome.
 	if a.OutputJSON {
  var summary *submit.Summary
  var ns, jobName string
@@ -843,6 +847,35 @@ other collaborators train against it without ever seeing the raw files.`))
  jsonEmitted = true
 	}
 
+	// Reclaim the staged source copy on a CLEAN success only. The
+	// ingestor copies (not moves) the staged files into the table, so
+	// leaving .tracebloc-staging/<table> behind doubles PVC use for
+	// file-bearing datasets until the next --overwrite or `data delete`
+	// (the staging-leak found by the ingest UX audit; cli#166 / epic #67).
+	// Gated on status=="succeeded" so we never touch the source on a
+	// - detached run (status "detached"): the Job is still reading it;
+	// - partial (status "completed_with_failures") or failed run: the
+	// user may want the source to inspect/retry.
+	// Best-effort and time-bounded (push.StagingCleanupTimeout): a failed
+	// or slow reclaim must not fail — or noticeably delay — a successful
+	// ingest.
+	if status == "succeeded" {
+ a.Printer.Infof("Reclaiming the temporary staging copy on the cluster…")
+ if cerr := push.CleanStaging(ctx, cs,
+ &push.SPDYExecutor{Config: resolved.RestConfig, Client: cs},
+ resolved.Namespace, a.Spec.Table, push.PodSpecOptions{
+ Namespace: resolved.Namespace,
+ PVCClaimName: pvc.ClaimName,
+ PVCMountPath: pvc.MountPath,
+ Table: a.Spec.Table,
+ ServiceAccountName: release.IngestorSAName,
+ Image: a.StagePodImage,
+ }); cerr != nil {
+ a.Printer.Warnf("Couldn't reclaim the temporary staging copy (%v). It's harmless — the next re-ingest of %q or a `tracebloc data delete %s` will clear it.",
+ cerr, a.Spec.Table, a.Spec.Table)
+ }
+	}
+
 	if exitErr != nil {
  return exitErr
 	}

internal/push/teardown.go

@@ -17,6 +17,13 @@ import (
 // and the teardown path agree on where a table lives.
 const IngestionDatabase = "training_test_datasets"
 
+// StagingCleanupTimeout bounds the best-effort post-success staging
+// reclaim (CleanStaging). The reclaim pod reuses the image the stage pod
+// just pulled, so it is normally Ready in seconds; this cap keeps a
+// stuck/unschedulable cleanup pod from adding the full pod-ready timeout
+// to a command the user already saw succeed.
+const StagingCleanupTimeout = 45 * time.Second
+
 // TeardownPlan enumerates the in-cluster artifacts `dataset rm` removes
 // for a pushed table: the MySQL table and the dataset's directories on
 // the shared PVC.
@@ -116,6 +123,65 @@ func Teardown(ctx context.Context, cs kubernetes.Interface, exec Executor, names
 	return res, nil
 }
 
+// CleanStaging best-effort removes ONLY the staged source copy at
+// StagedPrefix(table) from the shared PVC — never the final table dir
+// (FinalDestPrefix) and never the MySQL table.
+//
+// Why it's needed: the CLI streams a full copy of the dataset into
+// SharedRoot/.tracebloc-staging/<table>, and the in-cluster ingestor
+// COPIES (shutil.copy, not move) those files into the final table dir.
+// So after a successful load the staged source lingers on the PVC until
+// a later --overwrite or `data delete`, doubling disk use for
+// file-bearing datasets (image/detection/segmentation). Reclaiming it on
+// a clean success keeps the shared PVC from silently filling up.
+//
+// It reuses the same ephemeral stage-identity pod Teardown uses: that
+// pod runs as the uid that WROTE the staging files (65532), so it owns
+// them and the rm works by ownership on hostPath and CSI alike.
+//
+// Callers MUST treat a returned error as non-fatal — a leftover source
+// copy must never turn an otherwise-successful ingest into a failure —
+// and MUST only call this once the ingestion Job has SUCCEEDED (the
+// ingestor reads from this path while it runs; removing it mid-run, or
+// on a detached/failed run that may be retried, would corrupt the load).
+func CleanStaging(ctx context.Context, cs kubernetes.Interface, exec Executor, namespace, table string, podOpts PodSpecOptions) error {
+	// Panics on an unsafe name — callers stage only after ValidateTableName.
+	staged := StagedPrefix(table)
+
+	podOpts.Namespace = namespace
+	// Create under a detached, bounded context: a parent-ctx cancel
+	// (Ctrl-C) landing in the create window could otherwise drop a pod the
+	// apiserver already committed, orphaning it because the deferred delete
+	// below wouldn't yet have a name to reap. A fresh context keeps the
+	// create → deferred-delete pair atomic.
+	createCtx, cancelCreate := context.WithTimeout(context.Background(), StagingCleanupTimeout)
+	defer cancelCreate()
+	podName, err := CreateStagePod(createCtx, cs, podOpts)
+	if err != nil {
+ return fmt.Errorf("creating staging-cleanup pod: %w", err)
+	}
+	defer func() {
+ delCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+ defer cancel()
+ _ = DeleteStagePod(delCtx, cs, namespace, podName)
+	}()
+	// Bound the wait+exec so a stuck cleanup pod (unschedulable, slow image
+	// pull) can't tack the full 60s ready-timeout onto a command the user
+	// already saw succeed — while still honoring a parent-ctx cancel via the
+	// child. This reclaim is best-effort; the caller warns and moves on.
+	workCtx, cancelWork := context.WithTimeout(ctx, StagingCleanupTimeout)
+	defer cancelWork()
+	if _, err := WaitForStagePodReady(workCtx, cs, namespace, podName); err != nil {
+ return fmt.Errorf("waiting for staging-cleanup pod: %w", err)
+	}
+	var stderr bytes.Buffer
+	if err := exec.Exec(workCtx, namespace, podName, "stage",
+ []string{"rm", "-rf", staged}, nil, nil, &stderr); err != nil {
+ return fmt.Errorf("removing staged copy %s: %w%s", staged, err, stderrSuffix(&stderr))
+	}
+	return nil
+}
+
 // findRunningPod returns the name + first-container name of the first
 // Running pod in namespace whose name contains substr.
 func findRunningPod(ctx context.Context, cs kubernetes.Interface, namespace, substr string) (podName, container string, err error) {

internal/push/teardown_test.go

@@ -2,11 +2,13 @@ package push
 
 import (
 	"context"
+	"errors"
 	"strings"
 	"testing"
 
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/kubernetes/fake"
 	k8stesting "k8s.io/client-go/testing"
 )
@@ -115,3 +117,69 @@ func TestTeardown_RemovesViaStageIdentityPod(t *testing.T) {
  t.Errorf("Teardown leaked %d stage pod(s)", len(pods.Items))
 	}
 }
+
+// TestCleanStaging_RemovesOnlyStagingPrefix pins the staging-leak fix:
+// on a clean success the CLI reclaims ONLY .tracebloc-staging/<table>
+// (StagedPrefix) — never the final table dir (FinalDestPrefix) and never
+// the MySQL table — via the same ephemeral stage-identity pod Teardown
+// uses (so the rm works by ownership on hostPath + CSI).
+func TestCleanStaging_RemovesOnlyStagingPrefix(t *testing.T) {
+	cs := fake.NewClientset()
+	readyOnNextGet(cs)
+	fe := &fakeExecutor{}
+
+	if err := CleanStaging(context.Background(), cs, fe, "tracebloc", "reg_train", PodSpecOptions{
+ Namespace: "tracebloc",
+ PVCClaimName: "client-pvc",
+ PVCMountPath: "/data/shared",
+ Table: "reg_train",
+	}); err != nil {
+ t.Fatalf("CleanStaging: %v", err)
+	}
+
+	// The rm must target ONLY the staging prefix — not the final table dir.
+	wantCmd := "rm -rf " + StagedPrefix("reg_train")
+	if got := strings.Join(fe.gotCmd, " "); got != wantCmd {
+ t.Errorf("rm cmd = %q, want %q", got, wantCmd)
+	}
+	if strings.Contains(strings.Join(fe.gotCmd, " "), FinalDestPrefix("reg_train")) {
+ t.Errorf("rm cmd %q touched the final table dir — CleanStaging must never remove FinalDestPrefix", fe.gotCmd)
+	}
+
+	// It must run in the ephemeral stage-identity pod, not jobs-manager.
+	if !strings.HasPrefix(fe.gotPod, "tracebloc-stage-") {
+ t.Errorf("rm ran in pod %q, want the ephemeral stage pod (tracebloc-stage-*)", fe.gotPod)
+	}
+	if fe.gotContainer != "stage" {
+ t.Errorf("rm container = %q, want stage", fe.gotContainer)
+	}
+
+	// No leaked cleanup pods.
+	pods, _ := cs.CoreV1().Pods("tracebloc").List(context.Background(),
+ metav1.ListOptions{LabelSelector: StagePodManagedByLabel + "=" + StagePodManagedByValue})
+	if len(pods.Items) != 0 {
+ t.Errorf("CleanStaging leaked %d stage pod(s)", len(pods.Items))
+	}
+}
+
+// TestCleanStaging_PodCreateFailureReturnsError confirms the reclaim
+// surfaces a pod-create failure as an error (the caller logs it as a
+// non-fatal warning — a leftover staging copy must never fail an
+// otherwise-successful ingest).
+func TestCleanStaging_PodCreateFailureReturnsError(t *testing.T) {
+	cs := fake.NewClientset()
+	cs.PrependReactor("create", "pods", func(k8stesting.Action) (bool, runtime.Object, error) {
+ return true, nil, errors.New("PSA denied")
+	})
+	fe := &fakeExecutor{}
+
+	err := CleanStaging(context.Background(), cs, fe, "tracebloc", "reg_train", PodSpecOptions{
+ Namespace: "tracebloc", PVCClaimName: "client-pvc", PVCMountPath: "/data/shared", Table: "reg_train",
+	})
+	if err == nil {
+ t.Fatal("CleanStaging returned nil, want an error when the cleanup pod can't be created")
+	}
+	if fe.gotCmd != nil {
+ t.Errorf("rm ran (%v) despite the pod never being created", fe.gotCmd)
+	}
+}