发布

  • [OPIK-6186] [BE] feat: dataset project migration job (V1 → V2 auto-inference) (#6799)

    frostbyte_neo 发布于 2026-05-21 13:09:14 +00:00

    • [OPIK-6186] [BE] feat: add experiment and dataset migration skip columns to workspaces
    • Migration 000073: add experiment_project_migration_skipped_at /
      experiment_project_migration_skip_reason alongside the existing
      migration_skipped_at / migration_skipped_reason columns; copy any
      existing data into the new columns; keep old columns alive for
      backward compat during rolling deployment.

    • Migration 000074: add dataset_project_migration_skipped_at /
      dataset_project_migration_skip_reason columns (no legacy columns to
      mirror; these are net-new for the upcoming dataset migration job).

    • WorkspacesDAO / WorkspacesService: rename markMigrationSkipped /
      findMigrationSkippedWorkspaceIds / countMigrationSkipped to their
      experiment_project_migration_* equivalents; all UPDATE and INSERT
      statements dual-write to both old and new columns so old pods still
      reading migration_skipped_at see current data during the deployment
      window; add parallel dataset_project_migration_* methods.

    • Workspace record: replace generic migrationSkippedAt/Reason with
      experimentProjectMigrationSkippedAt/Reason +
      datasetProjectMigrationSkippedAt/Reason.

    • ExperimentProjectMigrationService + ExperimentProjectMigrationJobTest:
      update two call sites each to use the renamed service methods.

    • fix: read experiment migration skip status from legacy column during rolling deployment

    • fix: guard UPDATE on legacy column; add legacy fields to Workspace record

    • [OPIK-6186] [BE] feat: dataset project migration job (V1 → V2 auto-inference)

    Implements the dataset analogue of the experiment project migration (D1).
    Sets datasets.project_id from NULL by inferring the project from each
    dataset's experiments → experiment_items → traces graph in ClickHouse.

    Architecture (matches D1 sibling and the ticket's locked-in design):

    • Eligibility: pure MySQL on datasets (orphans excluding demo names and
      the env-var exclusion list, ordered smallest-first).
    • Inference: ClickHouse, joins experiments → experiment_items → traces,
      groups by dataset_id, exposes count(DISTINCT t.project_id) so the
      service does the four-bucket classification.
    • Four buckets — certain / certain-deleted / ambiguous / no-inference —
      classified in the service, not the SQL. No-inference orphans fall back
      to the workspace's "Default Project".
    • Three trap reasons persisted in workspaces.dataset_project_migration_*:
      deleted_project, all_ambiguous, default_project_missing. Also
      surfaced as labels on the cycle.trapped_workspaces gauge so dashboards
      can break down by reason.
    • Write: MySQL UPDATE with WHERE project_id IS NULL idempotency guard.
    • Post-migration: only workspaceVersionService.evictCache(workspaceId).

    Critical fix: when a workspace has both certain orphans and no-inference
    orphans without a Default Project, the certain mappings migrate BEFORE
    the workspace is trapped. Earlier draft trapped first and silently
    dropped the validated certain mappings.

    Query optimizations (validated against prod read replicas):

    • CH inference: per-alias workspace_id predicate on every joined alias.
      Without it the optimizer doesn't push the workspace filter through
      INNER JOIN, leading to a full traces-table scan that would exceed the
      cluster's row-read cap. With it, the primary key prunes most granules.
    • MySQL eligibility with exclusions: FORCE INDEX (datasets_workspace_id_name_uk) flips a full table scan into a range
      scan — about 4× faster end-to-end.
    • Secondary ORDER BY workspace_id ASC for deterministic tie-breaking
      (verified zero cost impact).

    Reactive flow:

    • All blocking JDBI calls wrapped in
      Mono.fromCallable(...).subscribeOn(migrationScheduler).
    • Dedicated bounded-elastic scheduler isolates JDBC work from the shared
      boundedElastic() and from the reactive client pools, mirroring D1.

    Metrics under opik.migration.dataset_project.*:

    • cycle.eligible_workspaces, cycle.trapped_workspaces{reason},
      cycle.env_excluded_workspaces, cycle.duration{result} (on Job),
      workspace.duration{result}, datasets.skipped{reason},
      datasets.assigned_to_default, batch.size.

    Tests: 10 integration tests in DatasetProjectMigrationJobTest covering
    every bucket, every trap reason, idempotency, env-excluded, and the
    critical-fix regression (certain mappings preserved when Default Project
    is missing). Container reuse disabled on this test class because the
    MySQL eligibility query is sensitive to stale orphan rows across runs.

    Collateral fixes (column-rename leftovers caught while running tests):

    • ExperimentProjectMigrationService.markMigrationSkipped(...)
      markExperimentProjectMigrationSkipped(...).
    • ExperimentProjectMigrationServiceTest: same call-site renames plus
      migrationSkippedReason()experimentProjectMigrationSkipReason().
    • fix(migration): address PR review feedback
    • WorkspacesDAO: change NULL guard on updateExperimentProjectMigrationSkippedIfNull
      from migration_skipped_at IS NULL to experiment_project_migration_skipped_at IS NULL.
      The legacy-column guard prevented new pods from backfilling rows written by old pods
      with only the legacy column set; the new column would have stayed NULL forever and
      the skip state would be lost when the legacy column gets dropped.

    • DatasetProjectMigrationJobTest: extend the 4 trap-path tests to assert the persisted
      datasetProjectMigrationSkipReason matches the expected catalog value, not just the
      presence of the workspace in the trapped list. Catches regressions where the trap
      fires but writes the wrong reason. Factored the assertion into a shared helper.

    • refactor(migration): align D2 with D1 policy (auto-create Default Project)

    Mirror the experiment-migration (D1) policy on the dataset job:

    • certain-deleted datasets reroute to Default Project instead of trapping
      the workspace; the lookup uses ProjectService.getOrCreate so a missing
      Default Project is auto-provisioned in-line.
    • no-inference datasets also flow through getOrCreate, eliminating the
      default_project_missing trap.
    • only all_ambiguous remains as an active trap reason — the CLI tool
      will cover those later.

    Tag the datasets.assigned_to_default counter with reason=deleted_project
    or no_inference so the dashboard can split the two contributing buckets.

    Drop dead code introduced by the removed trap paths: defaultProjectMissing
    parameter and unreachable branch, RESULT_DEFAULT_PROJECT_MISSING and
    RESULT_ALL_SKIPPED_DELETED attribute constants. Keep the trap-reason
    string constants in KNOWN_TRAP_REASONS so the gauge still buckets legacy
    workspaces persisted before the policy change.

    Fix a latent bug in resolveDefaultProjectAndMigrate: Mono.fromCallable
    returning null produces an empty Mono in Reactor, which short-circuited
    the happy path. Use an early return when no datasets need Default.

    Update affected tests: certain-deleted now migrates, default-missing now
    auto-creates, mixed-bucket workspace keeps only ambiguous as V1.

    • refactor(migration): address PR review feedback (OPIK-6186)

    Address all 13 unresolved review comments on PR #6799:

    Nits & cleanups:

    • MigrationSkipReasonCount: convert comment to javadoc, add @Builder and
      @NonNull (#1, #2).
    • DatasetDAO.batchSetProjectId: use text block instead of string
      concatenation (#7).
    • DatasetDAO.findEligibleDatasetMigrationWorkspaces: drop the default
      dispatch method and the no-exclusion SQL variant; collapse to one
      SqlQuery driven by <if(excludedWorkspaceIds)> so FORCE INDEX and
      workspace_id NOT IN (...) are emitted together exactly when needed
      (#4, #6). EXPLAIN against prod confirms: no exclusions → planner picks
      the right index on its own; with exclusions → FORCE INDEX is still
      required (otherwise type=ALL full scan, 220k rows).
    • DatasetProjectMigrationService: extract findEligibleWorkspaces,
      findOrphanDatasetIds, writeBatch private helpers around the
      transactionTemplate calls (#8).
    • config-test.yml: datasetProjectMigration.startupDelay 0s → 5s so the
      job test's seed completes before the first cycle fires (#12).
    • WorkspacesDAO/Service: remove unused countDatasetProjectMigrationSkipped
      (#3, the only DAO method that test scenarios needed but production
      doesn't).

    Substantive changes:

    • ExperimentDAO.computeDatasetProjectMapping: accept the orphan ID set
      and add e.dataset_id IN :dataset_ids so the join only walks the
      graph for V1 datasets (#10).
    • Same query rewritten to derive project_id from experiments.project_id
      (which D1 sets) instead of the experiments→experiment_items→traces
      join, with argMax(project_id, last_updated_at) GROUP BY id to dedup
      across ReplacingMergeTree row versions and HAVING != '' to filter
      experiments D1 left unmigrated (#11). Prod measurement on the
      worst-case workspace (120k experiments, 50-dataset sample):
      read_rows 1.15M → 130k (~8.8× fewer), read_bytes 127MB → 18MB
      (~7× less), latency 533ms → 111ms (~4.8× faster), and inference rate
      48% → 84% because experiment rows exist independently of trace links.

    Test split (#13):

    • New DatasetProjectMigrationServiceTest covers all classification and
      policy cases (9 tests) by calling runMigrationCycle().block()
      directly, mirroring ExperimentProjectMigrationServiceTest.
    • DatasetProjectMigrationJobTest trimmed to a single happy-path E2E
      driven by the scheduler. Combined suite: 10 tests in ~36s (was ~120s).

    Audit (no code change, will reply on the threads):

    • #5 datasets indexes: existing (workspace_id, name) UK and
      (workspace_id, project_id) idx cover both query patterns; orphan
      query plan = type=ref on workspace_project_idx.
    • #9 scheduler placement: every blocking JDBI / projectService /
      workspacesService call already has .subscribeOn(migrationScheduler);
      the R2DBC Flux from computeDatasetProjectMapping is non-blocking and
      the post-collect work is shifted to migrationScheduler via
      .publishOn().
    • [OPIK-6186] [BE] address PR review follow-ups on COMPUTE_DATASET_PROJECT_MAPPING

    Address the 4 unresolved threads on PR #6799:

    • ExperimentDAO COMPUTE_DATASET_PROJECT_MAPPING: drop redundant
      AND dataset_id != '' predicate (overlaps with dataset_id IN :dataset_ids and dataset_id is part of the PK).
    • ExperimentDAO.computeDatasetProjectMapping: switch the parameter from
      @NonNull Set<UUID> + isEmpty() to plain Set<UUID> +
      CollectionUtils.isEmpty() per reviewer preference.
    • ExperimentDAO row mapper: defensively wrap the project_id read in
      Optional.ofNullable(...).filter(StringUtils::isNotBlank) and emit
      via Mono::justOrEmpty, mirroring the experiment-side mapper on the
      same DAO. The SQL HAVING experiment_project_id != '' already filters
      blanks; this is belt-and-suspenders so a future schema change can't
      cause an UUID.fromString("") IAE. The service-level no-inference
      bucket already routes filtered-out datasets to the workspace's
      Default Project (or auto-provisions via projectService.getOrCreate).

    Tests:

    • DatasetProjectMigrationServiceTest: add two dedicated-workspace
      scenarios for the reviewer's edge case where every experiment for a
      dataset has project_id = '' (D1-pending or D1-ambiguous):
      • …ToExistingDefaultProject: pre-seeded Default Project; dataset
        migrates there.
      • …AutoCreatingDefaultProject: no Default Project; service
        auto-provisions one and migrates the dataset there.
        Both prove the empty-project-id rows don't slip through as
        certain/ambiguous and that the no-inference fallback covers them.
    下载附件