Skip to content

Commit

Permalink
Fix: Solved code quality issues
Browse files Browse the repository at this point in the history
  • Loading branch information
nesitor committed Dec 5, 2024
1 parent aaaec2a commit 24516a2
Show file tree
Hide file tree
Showing 8 changed files with 28 additions and 37 deletions.
5 changes: 0 additions & 5 deletions src/aleph/vm/controllers/firecracker/executable.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,11 +72,6 @@ class HostVolume:
read_only: bool


@dataclass
class HostGPU:
pci_host: str


@dataclass
class BaseConfiguration:
vm_hash: ItemHash
Expand Down
13 changes: 5 additions & 8 deletions src/aleph/vm/controllers/qemu/instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from asyncio import Task
from asyncio.subprocess import Process
from pathlib import Path
from typing import Generic, TypeVar, List
from typing import Generic, List, TypeVar

import psutil
from aleph_message.models import ItemHash
Expand All @@ -17,28 +17,28 @@
from aleph.vm.controllers.configuration import (
Configuration,
HypervisorType,
QemuGPU,
QemuVMConfiguration,
QemuVMHostVolume,
QemuGPU,
save_controller_configuration,
)
from aleph.vm.controllers.firecracker.executable import (
AlephFirecrackerResources,
VmSetupError,
HostGPU,
)
from aleph.vm.controllers.interface import AlephVmControllerInterface
from aleph.vm.controllers.qemu.cloudinit import CloudInitMixin
from aleph.vm.network.firewall import teardown_nftables_for_vm
from aleph.vm.network.interfaces import TapInterface
from aleph.vm.resources import HostGPU
from aleph.vm.storage import get_rootfs_base_path
from aleph.vm.utils import HostNotFoundError, ping, run_in_subprocess

logger = logging.getLogger(__name__)


class AlephQemuResources(AlephFirecrackerResources):
gpus: list[HostGPU]
gpus: List[HostGPU] = []

async def download_runtime(self) -> None:
volume = self.message_content.rootfs
Expand Down Expand Up @@ -204,10 +204,7 @@ async def configure(self):
)
for volume in self.resources.volumes
],
gpus=[
QemuGPU(pci_host=gpu.pci_host)
for gpu in self.resources.gpus
]
gpus=[QemuGPU(pci_host=gpu.pci_host) for gpu in self.resources.gpus],
)

configuration = Configuration(
Expand Down
5 changes: 1 addition & 4 deletions src/aleph/vm/controllers/qemu_confidential/instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,10 +127,7 @@ async def configure(self):
)
for volume in self.resources.volumes
],
gpus=[
QemuGPU(pci_host=gpu.pci_host)
for gpu in self.resources.gpus
]
gpus=[QemuGPU(pci_host=gpu.pci_host) for gpu in self.resources.gpus],
)

configuration = Configuration(
Expand Down
4 changes: 2 additions & 2 deletions src/aleph/vm/hypervisors/qemu/qemuvm.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import qmp
from systemd import journal

from aleph.vm.controllers.configuration import QemuVMConfiguration, QemuGPU
from aleph.vm.controllers.configuration import QemuGPU, QemuVMConfiguration
from aleph.vm.controllers.qemu.instance import logger


Expand Down Expand Up @@ -107,7 +107,7 @@ async def start(
# Use host-phys-bits-limit argument for GPU support. TODO: Investigate how to get the correct bits size
#
"-cpu",
"host,host-phys-bits-limit=0x28"
"host,host-phys-bits-limit=0x28",
# Uncomment for debug
# "-serial", "telnet:localhost:4321,server,nowait",
# "-snapshot", # Do not save anything to disk
Expand Down
2 changes: 1 addition & 1 deletion src/aleph/vm/hypervisors/qemu_confidential/qemuvm.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ async def start(
# AMD have different methods) and properly boot.
# Use host-phys-bits-limit argument for GPU support. TODO: Investigate how to get the correct bits size
"-cpu",
"host,host-phys-bits-limit=0x28"
"host,host-phys-bits-limit=0x28",
# Uncomment following for debug
# "-serial", "telnet:localhost:4321,server,nowait",
# "-snapshot", # Do not save anything to disk
Expand Down
14 changes: 6 additions & 8 deletions src/aleph/vm/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from aleph_message.models.execution.environment import HypervisorType

from aleph.vm.conf import settings
from aleph.vm.controllers.firecracker.executable import AlephFirecrackerExecutable, HostGPU
from aleph.vm.controllers.firecracker.executable import AlephFirecrackerExecutable
from aleph.vm.controllers.firecracker.instance import AlephInstanceResources
from aleph.vm.controllers.firecracker.program import (
AlephFirecrackerProgram,
Expand All @@ -38,7 +38,7 @@
)
from aleph.vm.orchestrator.pubsub import PubSub
from aleph.vm.orchestrator.vm import AlephFirecrackerInstance
from aleph.vm.resources import GpuDevice
from aleph.vm.resources import GpuDevice, HostGPU
from aleph.vm.systemd import SystemDManager
from aleph.vm.utils import create_task_log_exceptions, dumps_for_json

Expand Down Expand Up @@ -70,7 +70,9 @@ class VmExecution:
vm_hash: ItemHash
original: ExecutableContent
message: ExecutableContent
resources: AlephProgramResources | AlephInstanceResources | AlephQemuResources | AlephQemuConfidentialInstance | None = None
resources: (
AlephProgramResources | AlephInstanceResources | AlephQemuResources | AlephQemuConfidentialInstance | None
) = None
vm: AlephFirecrackerExecutable | AlephQemuInstance | AlephQemuConfidentialInstance | None = None
gpus: List[HostGPU]

Expand Down Expand Up @@ -224,11 +226,7 @@ def prepare_gpus(self, available_gpus: List[GpuDevice]) -> None:
for gpu in self.message.requirements.gpu:
for available_gpu in available_gpus:
if available_gpu.device_id == gpu.device_id:
gpus.append(
HostGPU(
pci_host=available_gpu.pci_host
)
)
gpus.append(HostGPU(pci_host=available_gpu.pci_host))
break
self.gpus = gpus

Expand Down
16 changes: 7 additions & 9 deletions src/aleph/vm/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from aleph.vm.controllers.firecracker.snapshot_manager import SnapshotManager
from aleph.vm.network.hostnetwork import Network, make_ipv6_allocator
from aleph.vm.orchestrator.metrics import get_execution_records
from aleph.vm.resources import GpuDevice, get_gpu_devices
from aleph.vm.resources import GpuDevice, HostGPU, get_gpu_devices
from aleph.vm.systemd import SystemDManager
from aleph.vm.utils import get_message_executable_content
from aleph.vm.vm_type import VmType
Expand Down Expand Up @@ -297,14 +297,12 @@ def get_instance_executions(self) -> Iterable[VmExecution]:
return executions or []

def get_available_gpus(self) -> List[GpuDevice]:
available_gpus = (
gpu
for gpu in self.gpus
for _, execution in self.executions.items()
if (isinstance(execution.resources, AlephQemuResources) or isinstance(execution.resources, AlephQemuConfidentialResources)) and not execution.uses_device_gpu(gpu.pci_host)
)

return available_gpus or []
available_gpus = []
for gpu in self.gpus:
for _, execution in self.executions.items():
if not execution.uses_gpu(gpu.pci_host):
available_gpus.append(gpu)
return available_gpus

def get_executions_by_sender(self, payment_type: PaymentType) -> dict[str, dict[str, list[VmExecution]]]:
"""Return all executions of the given type, grouped by sender and by chain."""
Expand Down
6 changes: 6 additions & 0 deletions src/aleph/vm/resources.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,17 @@
import subprocess
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional

from aleph_message.models import HashableModel
from pydantic import Extra, Field


@dataclass
class HostGPU:
pci_host: str


class GpuDeviceClass(str, Enum):
VGA_COMPATIBLE_CONTROLLER = "0300"
_3D_CONTROLLER = "0302"
Expand Down

0 comments on commit 24516a2

Please sign in to comment.