Skip to content

Commit

Permalink
fix: service creation failing
Browse files Browse the repository at this point in the history
  • Loading branch information
kyujin-cho committed May 22, 2024
1 parent 627fcba commit d82ba12
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 9 deletions.
23 changes: 19 additions & 4 deletions src/ai/backend/manager/api/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,11 @@ class RouteInfoModel(BaseModel):

class ServeInfoModel(BaseResponseModel):
endpoint_id: uuid.UUID = Field(description="Unique ID referencing the model service.")
model_id: Annotated[uuid.UUID, Field(description="ID of model VFolder.")]
extra_mounts: Annotated[
Sequence[uuid.UUID],
Field(description="List of extra VFolders which will be mounted to model service session."),
]
name: str = Field(description="Name of the model service.")
desired_session_count: NonNegativeInt = Field(
description="Number of identical inference sessions."
Expand Down Expand Up @@ -238,6 +243,8 @@ async def get_info(request: web.Request) -> ServeInfoModel:

return ServeInfoModel(
endpoint_id=endpoint.id,
model_id=endpoint.model,
extra_mounts=[m.vfid.folder_id for m in endpoint.extra_mounts],
name=endpoint.name,
desired_session_count=endpoint.desired_session_count,
active_routes=[
Expand Down Expand Up @@ -265,7 +272,13 @@ class ServiceConfigModel(BaseModel):
),
)

extra_mounts: Annotated[dict[uuid.UUID, MountOptionModel], Field(default={})]
extra_mounts: Annotated[
dict[uuid.UUID, MountOptionModel],
Field(
description="Specifications about extra VFolders mounted to model service session",
default={},
),
]

environ: dict[str, str] | None = Field(
description="Environment variables to be set inside the inference session",
Expand Down Expand Up @@ -485,14 +498,14 @@ async def create(request: web.Request, params: NewServiceRequestModel) -> ServeI
creation_config = params.config.model_dump()
creation_config["mounts"] = [
validation_result.model_id,
*[m.vfid for m in validation_result.extra_mounts],
*[m.vfid.folder_id for m in validation_result.extra_mounts],
]
creation_config["mount_map"] = {
validation_result.model_id: params.config.model_mount_destination,
**{m.vfid: m.kernel_path for m in validation_result.extra_mounts},
**{m.vfid.folder_id: m.kernel_path.as_posix() for m in validation_result.extra_mounts},
}
creation_config["mount_options"] = {
m.vfid: {"permission": m.mount_perm} for m in validation_result.extra_mounts
m.vfid.folder_id: {"permission": m.mount_perm} for m in validation_result.extra_mounts
}
sudo_session_enabled = request["user"]["sudo_session_enabled"]

Expand Down Expand Up @@ -565,6 +578,8 @@ async def create(request: web.Request, params: NewServiceRequestModel) -> ServeI

return ServeInfoModel(
endpoint_id=endpoint_id,
model_id=endpoint.model,
extra_mounts=[m.vfid.folder_id for m in endpoint.extra_mounts],
name=params.service_name,
desired_session_count=params.desired_session_count,
active_routes=[],
Expand Down
8 changes: 7 additions & 1 deletion src/ai/backend/manager/models/endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -530,7 +530,7 @@ async def check_extra_mounts(
"Same VFolder appears on both model specification and VFolder mount"
)

requested_mounts = [model_id, *extra_mounts.keys()]
requested_mounts = [*extra_mounts.keys()]
requested_mount_map: dict[str | uuid.UUID, str] = {
folder_id: options.mount_destination
for folder_id, options in extra_mounts.items()
Expand All @@ -543,6 +543,12 @@ async def check_extra_mounts(
}
for folder_id, options in extra_mounts.items()
}
log.debug(
"requested mounts: {}, mount_map: {}, mount_options: {}",
requested_mounts,
requested_mount_map,
requested_mount_options,
)
allowed_vfolder_types = await shared_config.get_vfolder_types()
vfolder_mounts = await prepare_vfolder_mounts(
conn,
Expand Down
10 changes: 7 additions & 3 deletions src/ai/backend/manager/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -3817,13 +3817,17 @@ async def handle_route_creation(
resource_policy,
SessionTypes.INFERENCE,
{
"mounts": [endpoint.model, *[m.vfid for m in endpoint.extra_mounts]],
"mounts": [endpoint.model, *[m.vfid.folder_id for m in endpoint.extra_mounts]],
"mount_map": {
endpoint.model: endpoint.model_mount_destination,
**{m.vfid: m.kernel_path for m in endpoint.extra_mounts},
**{
m.vfid.folder_id: m.kernel_path.as_posix()
for m in endpoint.extra_mounts
},
},
"mount_options": {
m.vfid: {"permission": m.mount_perm} for m in endpoint.extra_mounts
m.vfid.folder_id: {"permission": m.mount_perm}
for m in endpoint.extra_mounts
},
"environ": endpoint.environ,
"scaling_group": endpoint.resource_group,
Expand Down
5 changes: 4 additions & 1 deletion src/ai/backend/manager/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,10 @@ def __call__(self, lock_id: LockID, lifetime_hint: float) -> AbstractDistributed


class MountOptionModel(BaseModel):
mount_destination: Annotated[str | None, Field(default=None)]
mount_destination: Annotated[
str | None,
Field(description="Mount destination, defaults to /home/work/{folder_name}.", default=None),
]
type: Annotated[MountTypes, Field(default=MountTypes.BIND)]
permission: Annotated[
MountPermission | None,
Expand Down

0 comments on commit d82ba12

Please sign in to comment.