easy set-up for dummy agent's cpu core

lablup · Oct 25, 2024 · 5439ccb · 5439ccb
1 parent 66203e9
commit 5439ccb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/ai/backend/agent/dummy/config.py b/src/ai/backend/agent/dummy/config.py
@@ -7,7 +7,7 @@
 DEFAULT_CONFIG_PATH = Path.cwd() / "agent.dummy.toml"
 
 RandomRange = t.Tuple(t.ToFloat, t.ToFloat)
-core_idx = {0, 1, 2, 3, 4}
+num_core = 4
 
 
 dummy_local_config = t.Dict({
@@ -29,7 +29,7 @@
         }),
         t.Key("resource"): t.Dict({
             t.Key("cpu"): t.Dict({
-                t.Key("core-indexes", default=core_idx): tx.ToSet,
+                t.Key("num-core", default=num_core): t.ToInt,
             }),
             t.Key("memory"): t.Dict({
                 t.Key("size", default=34359738368): t.Int,

diff --git a/src/ai/backend/agent/dummy/intrinsic.py b/src/ai/backend/agent/dummy/intrinsic.py
@@ -79,7 +79,7 @@ def get_metadata(self) -> AcceleratorMetadata:
         }
 
     async def list_devices(self) -> Collection[AbstractComputeDevice]:
-        cores = self.resource_config["cpu"]["core-indexes"]
+        num_core: int = self.resource_config["cpu"]["num_core"]
         return [
             CPUDevice(
                 device_id=DeviceId(str(core_idx)),
@@ -88,7 +88,7 @@ async def list_devices(self) -> Collection[AbstractComputeDevice]:
                 memory_size=0,
                 processing_units=1,
             )
-            for core_idx in sorted(cores)
+            for core_idx in range(num_core)
         ]
 
     async def available_slots(self) -> Mapping[SlotName, Decimal]: