DevicePool Module Design Document

Overview

The DevicePool module is responsible for the asynchronous allocation and release management of devices such as Ascend, CUDA, and CPU. It uses a producer-consumer model to implement concurrent control of device resources.

Initialization Parameters

Parameter Name Type/Required Description
device_list List[int] (Optional) A list of device IDs (e.g., [0, 1]), defaults to [0]

Environment Variable Configuration

Supports setting the device list through the environment variable AIKG_DEVICES_LIST, formatted as comma-separated device IDs, e.g., 0,1,2

Class Method Descriptions

acquire_device()

Asynchronously acquires an available device.

Return Value

  • int: Device ID

release_device(device_id)

Asynchronously releases a specified device.

Parameters

  • device_id (int): The ID of the device to be released