Skip to content

Support image columns and media collection in multi-shard Slurm runs #1114

Open

Description

@nabinchha

Problem

Slurm planning explicitly rejects image columns when array_tasks.count > 1 in packages/data-designer-slurm/src/data_designer/slurm/planning/resolution.py.

Image generation can persist media separately from dataset rows. The current collection implementation merges verified Parquet outputs; it does not consolidate shard-local media files and preserve or rewrite their references. Removing the planning guard alone would therefore leave disk-backed image outputs unsupported.

Proposed solution

Define and implement image-output handling across shard generation, retries, winner publication, and collection. Include media artifacts in the output contract as needed, give them collision-safe identities, and ensure collected rows resolve to the correct images. Specify which persistence modes are supported.

Acceptance criteria

  • A dataset with an image column can run across multiple array tasks and be collected successfully.
  • Collected media references remain valid after collection, without depending on temporary allocation storage.
  • Media names cannot collide across shards or attempts.
  • Collection uses artifacts from winning attempts only and validates required files and their integrity.
  • Missing or corrupt media produces an actionable failure; retries and repeated collection do not publish partial output.
  • Tests cover supported persistence modes, shard retries, artifact references, and collection; documentation states any remaining restrictions.

Alternatives

Use a single array task for image generation, or assemble media and dataset references manually after separate runs.

Context

Related to #850. Plugin/custom-column sharding is tracked separately.

Activity

  1. self-assigned this
    on Oct 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions