async def shutdown_service(
self,
kernel_id, # type: str
service, # type: str
):
log.info('rpc::shutdown_service(k:{0}, app:{1})', kernel_id, service)
return await self.agent.shutdown_service(KernelId(UUID(kernel_id)), service)
async def execute(
self,
kernel_id, # type: str
api_version, # type: int
run_id, # type: str
mode, # type: Literal['query', 'batch', 'continue', 'input']
code, # type: str
opts, # type: Dict[str, Any]
flush_timeout, # type: float
):
# type: (...) -> Dict[str, Any]
if mode != 'continue':
log.info('rpc::execute(k:{0}, run-id:{1}, mode:{2}, code:{3!r})',
kernel_id, run_id, mode,
code[:20] + '...' if len(code) > 20 else code)
result = await self.agent.execute(
KernelId(UUID(kernel_id)),
run_id,
mode,
code,
opts=opts,
api_version=api_version,
flush_timeout=flush_timeout
)
return result
async def start_service(
self,
kernel_id, # type: str
service, # type: str
opts # type: Dict[str, Any]
):
# type: (...) -> Dict[str, Any]
log.info('rpc::start_service(k:{0}, app:{1})', kernel_id, service)
return await self.agent.start_service(KernelId(UUID(kernel_id)), service, opts)
async def execute_batch(
self,
kernel_id, # type: str
startup_command, # type: str
) -> None:
# TODO: use task-group to keep track of completion/cancellation
asyncio.create_task(self.agent.execute_batch(
KernelId(UUID(kernel_id)),
startup_command,
))
await asyncio.sleep(0)
async def restart_kernel(
self,
creation_id: str,
session_id: str,
kernel_id: str,
updated_config: dict,
):
log.info('rpc::restart_kernel(s:{0}, k:{1})', session_id, kernel_id)
return await self.agent.restart_kernel(
creation_id,
SessionId(UUID(session_id)),
KernelId(UUID(kernel_id)),
cast(KernelCreationConfig, updated_config),
)
async def get_kernel_id_from_container(
val: Union[str, DockerContainer]) -> Optional[KernelId]:
if isinstance(val, DockerContainer):
if 'Name' not in val._container:
await val.show()
name = val['Name']
elif isinstance(val, str):
name = val
name = name.lstrip('/')
if not name.startswith('kernel.'):
return None
try:
return KernelId(UUID(name.rsplit('.', 2)[-1]))
except (IndexError, ValueError):
return None
async def destroy_kernel(
self,
kernel_id: str,
reason: str = None,
suppress_events: bool = False,
):
log.info('rpc::destroy_kernel(k:{0})', kernel_id)
done = asyncio.Event()
await self.agent.inject_container_lifecycle_event(
KernelId(UUID(kernel_id)),
LifecycleEvent.DESTROY,
reason or 'user-requested',
done_event=done,
suppress_events=suppress_events,
)
await done.wait()
return getattr(done, '_result', None)
async def create_kernels(
self,
creation_id: str,
raw_session_id: str,
raw_kernel_ids: Sequence[str],
raw_configs: Sequence[dict],
raw_cluster_info: dict,
):
cluster_info = cast(ClusterInfo, raw_cluster_info)
session_id = SessionId(UUID(raw_session_id))
raw_results = []
coros = []
for raw_kernel_id, raw_config in zip(raw_kernel_ids, raw_configs):
log.info('rpc::create_kernel(k:{0}, img:{1})',
raw_kernel_id, raw_config['image']['canonical'])
kernel_id = KernelId(UUID(raw_kernel_id))
kernel_config = cast(KernelCreationConfig, raw_config)
coros.append(self.agent.create_kernel(
creation_id,
session_id,
kernel_id,
kernel_config,
cluster_info,
))
results = await asyncio.gather(*coros, return_exceptions=True)
errors = [*filter(lambda item: isinstance(item, Exception), results)]
if errors:
# Raise up the first error.
raise errors[0]
raw_results = [
{
'id': str(result['id']),
'kernel_host': result['kernel_host'],
'repl_in_port': result['repl_in_port'],
'repl_out_port': result['repl_out_port'],
'stdin_port': result['stdin_port'], # legacy
'stdout_port': result['stdout_port'], # legacy
'service_ports': result['service_ports'],
'container_id': result['container_id'],
'resource_spec': result['resource_spec'],
'attached_devices': result['attached_devices'],
}
for result in results
]
return raw_results
'rocm.devices': Decimal('8'),
}),
),
]
@attr.s(auto_attribs=True, slots=True)
class SessionKernelIdPair:
session_id: UUID
kernel_ids: Sequence[KernelId]
pending_session_kernel_ids = [
SessionKernelIdPair(
session_id=UUID('251907d9-1290-4126-bc6c-000000000100'),
kernel_ids=[KernelId(UUID('251907d9-1290-4126-bc6c-000000000100'))]),
SessionKernelIdPair(
session_id=UUID('251907d9-1290-4126-bc6c-000000000200'),
kernel_ids=[KernelId(UUID('251907d9-1290-4126-bc6c-000000000200'))]),
SessionKernelIdPair(
# single-node mode multi-container session
session_id=UUID('251907d9-1290-4126-bc6c-000000000300'),
kernel_ids=[
KernelId(UUID('251907d9-1290-4126-bc6c-000000000300')),
KernelId(UUID('251907d9-1290-4126-bc6c-000000000301')),
KernelId(UUID('251907d9-1290-4126-bc6c-000000000302')),
]),
]
existing_session_kernel_ids = [
SessionKernelIdPair(
async def list_files(self, kernel_id: str, path: str):
log.info('rpc::list_files(k:{0}, fn:{1})', kernel_id, path)
return await self.agent.list_files(KernelId(UUID(kernel_id)), path)
async def download_file(self, kernel_id: str, filepath: str):
log.info('rpc::download_file(k:{0}, fn:{1})', kernel_id, filepath)
return await self.agent.download_file(KernelId(UUID(kernel_id)), filepath)
async def upload_file(self, kernel_id: str, filename: str, filedata: bytes):
log.info('rpc::upload_file(k:{0}, fn:{1})', kernel_id, filename)
await self.agent.accept_file(KernelId(UUID(kernel_id)), filename, filedata)
async def get_logs(self, kernel_id: str):
log.info('rpc::get_logs(k:{0})', kernel_id)
return await self.agent.get_logs(KernelId(UUID(kernel_id)))
async def get_completions(self, kernel_id: str,
text: str, opts: dict):
log.debug('rpc::get_completions(k:{0}, ...)', kernel_id)
await self.agent.get_completions(KernelId(UUID(kernel_id)), text, opts)
async def interrupt_kernel(self, kernel_id: str):
log.info('rpc::interrupt_kernel(k:{0})', kernel_id)
await self.agent.interrupt_kernel(KernelId(UUID(kernel_id)))
'mem': Decimal('0'),
'cuda.shares': Decimal('0'),
'rocm.devices': Decimal('0'),
}),
occupied_slots=ResourceSlot({
'cpu': Decimal('3.0'),
'mem': Decimal('2560'),
'cuda.shares': Decimal('1.0'),
'rocm.devices': Decimal('8'),
}),
),
]
pending_kernel_ids: Sequence[KernelId] = [
KernelId(uuid.uuid4()) for _ in range(3)
]
existing_kernel_ids: Sequence[KernelId] = [
KernelId(uuid.uuid4()) for _ in range(3)
]
_common_dummy_for_pending_session: Mapping[str, Any] = dict(
image_ref=ImageRef('lablup/python:3.6-ubunt18.04'),
domain_name='default',
group_id=example_group_id,
resource_policy={},
resource_opts={},
mounts=[],
mount_map={},
environ={},