|
3 | 3 | from functools import lru_cache |
4 | 4 | from typing import Optional |
5 | 5 |
|
6 | | -import cpuinfo |
7 | 6 | import psutil |
8 | 7 | from aiohttp import web |
| 8 | +from aleph.vm.conf import settings |
| 9 | +from aleph.vm.orchestrator.machine import get_cpu_info, get_memory_info |
9 | 10 | from aleph_message.models import ItemHash |
10 | 11 | from aleph_message.models.execution.environment import CpuProperties |
11 | 12 | from pydantic import BaseModel, Field |
12 | 13 |
|
13 | | -from aleph.vm.conf import settings |
14 | | - |
15 | 14 |
|
16 | 15 | class Period(BaseModel): |
17 | 16 | datetime: datetime |
@@ -76,18 +75,63 @@ class MachineUsage(BaseModel): |
76 | 75 | active: bool = True |
77 | 76 |
|
78 | 77 |
|
| 78 | +class ExtendedCpuProperties(CpuProperties): |
| 79 | + """CPU properties.""" |
| 80 | + |
| 81 | + model: Optional[str] = Field(default=None, description="CPU model") |
| 82 | + frequency: Optional[str] = Field(default=None, description="CPU frequency") |
| 83 | + |
| 84 | + |
| 85 | +class MemoryProperties(BaseModel): |
| 86 | + """MEMORY properties.""" |
| 87 | + |
| 88 | + size: Optional[str] = Field(default=None, description="Memory size") |
| 89 | + units: Optional[str] = Field(default=None, description="Memory size units") |
| 90 | + type: Optional[str] = Field(default=None, description="Memory type") |
| 91 | + clock: Optional[str] = Field(default=None, description="Memory clock") |
| 92 | + clock_units: Optional[str] = Field(default=None, description="Memory clock units") |
| 93 | + |
| 94 | + |
| 95 | +class MachineCapability(BaseModel): |
| 96 | + cpu: ExtendedCpuProperties |
| 97 | + memory: MemoryProperties |
| 98 | + |
| 99 | + |
79 | 100 | @lru_cache |
80 | 101 | def get_machine_properties() -> MachineProperties: |
81 | 102 | """Fetch machine properties such as architecture, CPU vendor, ... |
82 | 103 | These should not change while the supervisor is running. |
83 | 104 |
|
84 | 105 | In the future, some properties may have to be fetched from within a VM. |
85 | 106 | """ |
86 | | - cpu_info = cpuinfo.get_cpu_info() # Slow |
| 107 | + |
| 108 | + cpu_info = get_cpu_info() |
87 | 109 | return MachineProperties( |
88 | 110 | cpu=CpuProperties( |
89 | | - architecture=cpu_info["raw_arch_string"], |
90 | | - vendor=cpu_info["vendor_id"], |
| 111 | + architecture=cpu_info["architecture"], |
| 112 | + vendor=cpu_info["vendor"], |
| 113 | + ), |
| 114 | + ) |
| 115 | + |
| 116 | + |
| 117 | +@lru_cache |
| 118 | +def get_machine_capability() -> MachineCapability: |
| 119 | + cpu_info = get_cpu_info() |
| 120 | + mem_info = get_memory_info() |
| 121 | + |
| 122 | + return MachineCapability( |
| 123 | + cpu=ExtendedCpuProperties( |
| 124 | + architecture=cpu_info["architecture"], |
| 125 | + vendor=cpu_info["vendor"], |
| 126 | + model=cpu_info["model"], |
| 127 | + frequency=cpu_info["frequency"], |
| 128 | + ), |
| 129 | + memory=MemoryProperties( |
| 130 | + size=mem_info["size"], |
| 131 | + units=mem_info["units"], |
| 132 | + type=mem_info["type"], |
| 133 | + clock=mem_info["clock"], |
| 134 | + clock_units=mem_info["clock_units"], |
91 | 135 | ), |
92 | 136 | ) |
93 | 137 |
|
@@ -119,6 +163,13 @@ async def about_system_usage(_: web.Request): |
119 | 163 | return web.json_response(text=usage.json(exclude_none=True), headers={"Access-Control-Allow-Origin:": "*"}) |
120 | 164 |
|
121 | 165 |
|
| 166 | +async def about_capability(_: web.Request): |
| 167 | + """Public endpoint to expose information about the CRN capability.""" |
| 168 | + |
| 169 | + capability: MachineCapability = get_machine_capability() |
| 170 | + return web.json_response(text=capability.json(exclude_none=False), headers={"Access-Control-Allow-Origin:": "*"}) |
| 171 | + |
| 172 | + |
122 | 173 | class Allocation(BaseModel): |
123 | 174 | """An allocation is the set of resources that are currently allocated on this orchestrator. |
124 | 175 | It contains the item_hashes of all persistent VMs, instances, on-demand VMs and jobs. |
|
0 commit comments