Skip to content

Commit 7bf21b0

Browse files
nesitorolethanh
authored andcommitted
Fix: Solve message class casting issue with requirements field.
1 parent 28582cf commit 7bf21b0

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

src/aleph/vm/pool.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ async def create_a_vm(
149149
try:
150150
if message.requirements and message.requirements.gpu:
151151
# Ensure we have the necessary GPU for the user by reserving them
152-
resources = self.find_resources_available_for_user(cast(message, InstanceContent), message.address)
152+
resources = self.find_resources_available_for_user(message, message.address)
153153
# First assign Host GPUs from the available
154154
execution.prepare_gpus(list(resources))
155155
# Prepare VM general Resources and also the GPUs
@@ -383,7 +383,7 @@ def get_valid_reservation(self, resource) -> Reservation | None:
383383
del self.reservations[resource]
384384
return self.reservations.get(resource)
385385

386-
async def reserve_resources(self, message: InstanceContent, user):
386+
async def reserve_resources(self, message: ExecutableContent, user):
387387
gpu_to_reserve = message.requirements.gpu if message.requirements and message.requirements.gpu else []
388388
expiration_date = datetime.now(tz=timezone.utc) + timedelta(seconds=60)
389389
if not gpu_to_reserve:
@@ -400,8 +400,8 @@ async def reserve_resources(self, message: InstanceContent, user):
400400

401401
return expiration_date
402402

403-
def find_resources_available_for_user(self, message: InstanceContent, user) -> set[GpuDevice]:
404-
"""Find required resource to run InstanceContent from reserved resources by user or free resources.
403+
def find_resources_available_for_user(self, message: ExecutableContent, user) -> set[GpuDevice]:
404+
"""Find required resource to run ExecutableContent from reserved resources by user or free resources.
405405
406406
Only implement GPU for now"""
407407
# Calling function should use the creation_lock to avoid resource being stollem

0 commit comments

Comments
 (0)