How to set resources in pipeline component in Azure machine learning sdk v2?

Question

I'm creating a pipeline that scores images using a custom pytorch unet. Assume I registered my model and now I'm creating that pipeline. I'm running into a OOM issue due to Azure setting the docker shared memory to 2g.I attempted to set the resources in my component yml as:

schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
name: my_component
display_name: Create segmentation maps
type: command
description:  |-
  **U NNET**:

inputs:
  nnetmodel: 
    type: custom_model 
  prepared_data:
    type: uri_folder
outputs:
  score_results:
    type: uri_folder
  score_summary:
    type: uri_file
code: src
environment: azureml:mynet-env@latest
distribution:
  type: pytorch
  process_count_per_instance: 1
resources:
  shm_size: 20g
  instance_count: 1
  docker_args: --ipc=host --shm-size 20g
command: >-

Load component all looks good.

score_data = load_component(source="my_component.yml")
print(score_data)

But when I create/update the component the component in my main code it removes all the resources BUT the instance_count.

ml_client.components.create_or_update(score_data)

How can I change the docker configurations on Azure ML SDK V2 for pipeline components?

I expected ml_client.components.create_or_update(score_data) to keep the properties of the component I created.

Answer 1

It looks like this is a feature in Azure ML. I.E. whe a component is created/updated it will remove most of the settings. The ONLY resource setting that would stay in place is the instance_count.

Thus the correct order of steps to preserve the resources configuration is:

1. load_component

process_data = load_component(source="components/my_component/my_component.yml")

2. use the component in a new pipeline while creating your pipeline

    @pipeline()
    def mypipe(input_path: Input(type=AssetTypes.URI_FOLDER)):
        output_data = process_data(batch_path=input_path)
        return {"output_data ": output_data.outputs.score_results}

You can stop here if this is a single-use scenario or if you don't want to reuse the same pipeline.

To reuse this same pipeline but keeping the resources configuration you can register it in a batch endpoint.

deployment = PipelineComponentBatchDeployment(
    name="my-component-pipe",
    description="Keep resources :)",
    endpoint_name=endpoint.name, # This assumes you have created the endpoint.
    component=mypipe().component,
    settings={"continue_on_step_failure": False, "default_compute": compute_target},
)
# NOW you can register your deployment to be used.
ml_client.batch_deployments.begin_create_or_update(deployment).result()

Enjoy and let me know if there is something I missed.