diff --git a/cyborg/accelerator/drivers/aichip/__init__.py b/cyborg/accelerator/drivers/aichip/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/cyborg/accelerator/drivers/aichip/huawei/__init__.py b/cyborg/accelerator/drivers/aichip/huawei/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/cyborg/accelerator/drivers/aichip/huawei/ascend.py b/cyborg/accelerator/drivers/aichip/huawei/ascend.py new file mode 100644 index 00000000..9faea072 --- /dev/null +++ b/cyborg/accelerator/drivers/aichip/huawei/ascend.py @@ -0,0 +1,119 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +from cyborg.accelerator.drivers.driver import GenericDriver +from cyborg.objects.driver_objects import driver_deployable, driver_device, \ + driver_attach_handle, driver_controlpath_id + +import re +import subprocess + +from cyborg.accelerator.common import utils +from cyborg.common import constants + +from oslo_serialization import jsonutils + +PCI_INFO_PATTERN = re.compile("(?P[0-9a-f]{4}:[0-9a-f]{2}:" + "[0-9a-f]{2}\.[0-9a-f]) " + "(?P.*) [\[].*]: (?P.*) .*" + "[\[](?P[0-9a-fA-F]" + "{4}):(?P[0-9a-fA-F]{4})].*" + "[(rev ](?P[0-9a-f]{2})") + + +class AscendDriver(GenericDriver): + """The class for Ascend AI Chip drivers. + + This is the Huawei Ascend AI Chip drivers. + """ + VENDOR = "huawei" + + # TODO(yikun): can be extracted into PCIDeviceDriver + def _generate_controlpath_id(self, pci): + driver_cpid = driver_controlpath_id.DriverControlPathID() + driver_cpid.cpid_type = "PCI" + driver_cpid.cpid_info = pci["slot_json"] + return driver_cpid + + # TODO(yikun): can be extracted into PCIDeviceDriver + def _generate_attach_handle(self, pci): + driver_ah = driver_attach_handle.DriverAttachHandle() + driver_ah.attach_type = "PCI" + driver_ah.in_use = False + driver_ah.attach_info = pci["slot_json"] + return driver_ah + + # TODO(yikun): can be extracted into PCIDeviceDriver + def _generate_dep_list(self, pci): + driver_dep = driver_deployable.DriverDeployable() + driver_dep.attach_handle_list = [self._generate_attach_handle(pci)] + pci_addr_name = pci["slot"].replace(":", "_").replace(".", "_") + driver_dep.name = pci.get('device', '') + '_' + pci_addr_name + driver_dep.num_accelerators = 1 + driver_dep.driver_name = self.VENDOR + return [driver_dep] + + # TODO(yikun): can be extracted into PCIDeviceDriver + def _get_pci_lines(self, keywords=()): + cmd = "sudo lspci -nnn -D" + if keywords: + cmd += "| grep -E %s" % '|'.join(keywords) + # FIXME(wangzhh): Use oslo.privsep instead of subprocess here to + # prevent shell injection attacks. + p = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True) + p.wait() + pci_lines = p.stdout.readlines() + return pci_lines + + def discover(self): + """ + The PCI line would be matched as: + + 0000:0c:00.0 Processing acc [1200]: Device [19e5:d100] (rev 20) + + { + 'slot': '0000:0c:00.0', # domain:bus:device.function + 'device': 'Device', # Name of the device + 'vendor_id': '19e5', # ID of the vendor + 'class': 'Processing accelerators', # Name of the class + 'device_id': 'd100', # ID of the device + 'revision': '20' # Revision number + } + """ + ascends = self._get_pci_lines(('d100',)) + npu_list = [] + for ascend in ascends: + m = PCI_INFO_PATTERN.match(ascend) + if m: + pci_dict = m.groupdict() + pci_dict["slot_json"] = utils.pci_str_to_json(pci_dict["slot"]) + device = driver_device.DriverDevice() + device.stub = False + device.vendor = pci_dict["vendor_id"] + device.model = pci_dict.get('model', '') + std_board_info = {'device_id': pci_dict.get('device_id', None), + 'class': pci_dict.get('class', None)} + device.std_board_info = jsonutils.dumps(std_board_info) + device.vendor_board_info = '' + device.type = constants.DEVICE_AICHIP + device.controlpath_id = self._generate_controlpath_id(pci_dict) + device.deployable_list = self._generate_dep_list(pci_dict) + npu_list.append(device) + return npu_list + + def update(self, control_path, image_path): + # TODO(yikun): To be implemented in future + return True + + def get_stats(self): + # TODO(yikun): To be implemented in future + return {} diff --git a/cyborg/tests/unit/accelerator/drivers/aichip/__init__.py b/cyborg/tests/unit/accelerator/drivers/aichip/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/cyborg/tests/unit/accelerator/drivers/aichip/huawei/__init__.py b/cyborg/tests/unit/accelerator/drivers/aichip/huawei/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/cyborg/tests/unit/accelerator/drivers/aichip/huawei/test_ascend.py b/cyborg/tests/unit/accelerator/drivers/aichip/huawei/test_ascend.py new file mode 100644 index 00000000..a000b9d5 --- /dev/null +++ b/cyborg/tests/unit/accelerator/drivers/aichip/huawei/test_ascend.py @@ -0,0 +1,54 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +import json +import mock + +from cyborg.tests import base +from cyborg.accelerator.drivers.aichip.huawei.ascend import AscendDriver + +d100_pci_res = [ + "0000:00:0c.0 Processing accelerators [1200]:" + " Device [19e5:d100] (rev 20)\n", + "0000:00:0d.0 Processing accelerators [1200]:" + " Device [19e5:d100] (rev 20)" +] + + +class TestAscendDriver(base.TestCase): + + @mock.patch('cyborg.accelerator.drivers.aichip.' + 'huawei.ascend.AscendDriver._get_pci_lines', + return_value=d100_pci_res) + def test_discover(self, mock_pci): + ascend_driver = AscendDriver() + npu_list = ascend_driver.discover() + self.assertEqual(2, len(npu_list)) + for ascend in npu_list: + self.assertEqual('AICHIP', ascend.type) + self.assertEqual('PCI', ascend.controlpath_id.cpid_type) + self.assertEqual(json.loads( + '{"class": "Processing accelerators", "device_id": "d100"}'), + json.loads(ascend.std_board_info)) + self.assertEqual('19e5', ascend.vendor) + + self.assertEqual( + {"device": "0c", "bus": "00", "domain": "0000", "function": "0"}, + json.loads(npu_list[0].controlpath_id.cpid_info)) + self.assertEqual( + {"device": "0d", "bus": "00", "domain": "0000", "function": "0"}, + json.loads(npu_list[1].controlpath_id.cpid_info)) + + self.assertEqual('Device_0000_00_0c_0', + npu_list[0].deployable_list[0].name) + self.assertEqual('Device_0000_00_0d_0', + npu_list[1].deployable_list[0].name) diff --git a/setup.cfg b/setup.cfg index b9716dc1..7e815731 100644 --- a/setup.cfg +++ b/setup.cfg @@ -49,6 +49,7 @@ cyborg.accelerator.driver = nvmf_spdk_driver = cyborg.accelerator.drivers.spdk.nvmf.nvmf:NVMFDRIVER nvidia_gpu_driver = cyborg.accelerator.drivers.gpu.nvidia.driver:NVIDIAGPUDriver fake_driver = cyborg.accelerator.drivers.fake:FakeDriver + huawei_ascend_driver = cyborg.accelerator.drivers.aichip.huawei.ascend:AscendDriver oslo.config.opts = cyborg = cyborg.conf.opts:list_opts