Azure ML SDKv1：分区TabularDataset传入ParallelRunStep报错排查

问题现象 #

运行ParallelRunStep时触发以下错误：

UserInputNotPartitionedByGivenKeys: The input dataset 'partitioned_combined_scored_dataset_input' is not partitioned by 'model_name'.

完整报错堆栈：

Traceback (most recent call last):
  File "/tmp/48a0ec47-b89c-41ff-89f8-3482d2823d20/prs_prod/lib/python3.8/site-packages/azureml_sys/parallel_run/master_role_process.py", line 111, in run
    loop.run_until_complete(self.master_role.start())
  File "/tmp/48a0ec47-b89c-41ff-89f8-3482d2823d20/prs_prod/lib/python3.8/asyncio/base_events.py", line 616, in run_until_complete
    return future.result()
  File "/tmp/48a0ec47-b89c-41ff-89f8-3482d2823d20/prs_prod/lib/python3.8/site-packages/azureml_sys/parallel_run/master_role.py", line 303, in start
    await self.wait_for_first_task()
  File "/tmp/48a0ec47-b89c-41ff-89f8-3482d2823d20/prs_prod/lib/python3.8/site-packages/azureml_sys/parallel_run/master_role.py", line 288, in wait_for_first_task
    await self.wait_for_input_init()
  File "/tmp/48a0ec47-b89c-41ff-89f8-3482d2823d20/prs_prod/lib/python3.8/site-packages/azureml_sys/parallel_run/master_role.py", line 126, in wait_for_input_init
    self.future_create_tasks.result()
  File "/tmp/48a0ec47-b89c-41ff-89f8-3482d2823d20/prs_prod/lib/python3.8/site-packages/azureml_sys/parallel_run/task_producer.py", line 199, in create_tasks
    raise exc
  File "/tmp/48a0ec47-b89c-41ff-89f8-3482d2823d20/prs_prod/lib/python3.8/site-packages/azureml_sys/parallel_run/task_producer.py", line 190, in create_tasks
    for task_group in self.get_task_groups(provider.get_tasks()):
  File "/tmp/48a0ec47-b89c-41ff-89f8-3482d2823d20/prs_prod/lib/python3.8/site-packages/azureml_sys/parallel_run/task_producer.py", line 169, in get_task_groups
    for index, task in enumerate(tasks):
  File "/tmp/48a0ec47-b89c-41ff-89f8-3482d2823d20/prs_prod/lib/python3.8/site-packages/azureml_sys/parallel_run/partition_by_keys_provider.py", line 77, in get_tasks
    raise UserInputNotPartitionedByGivenKeys(message=message, compliant_message=compliant_message)
UserInputNotPartitionedByGivenKeys: The input dataset 'partitioned_combined_scored_dataset_input' is not partitioned by 'model_name'.

相关代码片段：

ParallelRunConfig与ParallelRunStep定义 #

parallel_run_config = ParallelRunConfig(
    source_directory=source_dir_for_snapshot,
    entry_script="src/steps/script.py",
    partition_keys=["model_name"], 
    error_threshold=10,        
    allowed_failed_count=15,
    allowed_failed_percent=10,
    run_max_try=3,
    output_action="append_row",
    append_row_file_name="output_file.csv",
    environment=aml_run_config.environment,
    compute_target=aml_run_config.target,
    node_count=2
)

parallelrun_step = ParallelRunStep(
    name="Do Some Parallel Stuff on Each model_name",
    parallel_run_config=parallel_run_config ,
    inputs=[partitioned_combined_scored_dataset],
    output=OutputFileDatasetConsumptionConfig(name='output_dataset'),
    arguments=["--score-id", score_id_pipeline_param,
               "--partitioned-combined-dataset", partitioned_combined_scored_dataset],
    allow_reuse=True
)

数据集配置与创建 #

# 数据集消费配置
partitioned_combined_scored_dataset = DatasetConsumptionConfig(
    name="partitioned_combined_scored_dataset_input",
    dataset=PipelineParameter(
        name="partitioned_combined_dataset",
        default_value=future_partitioned_dataset)
)

# 初始分区数据集创建
partitioned_dataset = TabularDatasetFactory.from_parquet_files(path=(Datastore.get(ws, ), f"{partitioned_combined_datasets_dir}/*.parquet"))\
    .partition_by(
        partition_keys=['model_name'], 
        target=DataPath(Datastore(), 'some/path/to/partitioned')
    )

# 运行时传入的数据集创建（导致问题的代码）
TabularDatasetFactory.from_parquet_files(path=(Datastore(), f"{partitioned_combined_dataset_dir}/*/*/*.parquet"))

错误原因 #

运行时通过通配符*/*/*.parquet创建的TabularDataset没有保留分区元数据。ParallelRunStep要求输入数据集必须是通过partition_by()方法生成、或显式指定了partition_keys的数据集——仅文件路径符合分区结构并不足以让AML识别其为分区数据集。

解决方案 #

方案1：直接使用partition_by()生成的数据集 #

无需在运行时重新创建数据集，直接将partition_by()返回的partitioned_dataset作为PipelineParameter的默认值或运行时传入值，它已经包含了model_name的分区元数据，AML可以直接识别。

方案2：运行时创建数据集时显式添加分区元数据 #

如果必须动态指定路径，使用with_partition_keys()方法为数据集添加分区键信息，确保AML能识别分区结构：

# 修改运行时的数据集创建代码
runtime_dataset = TabularDatasetFactory.from_parquet_files(
    path=(Datastore(), f"{partitioned_combined_dataset_dir}/*/*/*.parquet")
).with_partition_keys(['model_name'])

注意：需确保文件路径的层级与指定的分区键对应，AML会自动从路径中提取model_name的分区值。

方案3：检查PipelineParameter传值逻辑 #

确认运行时传入PipelineParameter的是带有分区元数据的TabularDataset对象，而非仅读取文件的普通数据集。

内容的提问来源于stack exchange，提问作者geominded