More efficiency (#12112)

This commit is contained in:
Tim O'Farrell 2025-12-19 21:18:48 -07:00 committed by GitHub
parent fa2567b2a0
commit 0677cebb25
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 39 additions and 18 deletions

View File

@ -252,10 +252,9 @@ class RemoteSandboxService(SandboxService):
# The batch endpoint should return a list of runtimes
# Convert to a dictionary keyed by session_id for easy lookup
runtimes_by_id = {}
if batch_data and 'runtimes' in batch_data:
for runtime in batch_data['runtimes']:
if 'session_id' in runtime:
runtimes_by_id[runtime['session_id']] = runtime
for runtime in batch_data:
if runtime and 'session_id' in runtime:
runtimes_by_id[runtime['session_id']] = runtime
return runtimes_by_id
@ -566,6 +565,32 @@ class RemoteSandboxService(SandboxService):
return paused_sandbox_ids
async def batch_get_sandboxes(
self, sandbox_ids: list[str]
) -> list[SandboxInfo | None]:
"""Get a batch of sandboxes, returning None for any which were not found."""
if not sandbox_ids:
return []
query = await self._secure_select()
query = query.filter(StoredRemoteSandbox.id.in_(sandbox_ids))
stored_remote_sandboxes = await self.db_session.execute(query)
stored_remote_sandboxes_by_id = {
stored_remote_sandbox[0].id: stored_remote_sandbox[0]
for stored_remote_sandbox in stored_remote_sandboxes
}
runtimes_by_id = await self._get_runtimes_batch(
list(stored_remote_sandboxes_by_id)
)
results = []
for sandbox_id in sandbox_ids:
stored_remote_sandbox = stored_remote_sandboxes_by_id.get(sandbox_id)
result = None
if stored_remote_sandbox:
runtime = runtimes_by_id.get(sandbox_id)
result = self._to_sandbox_info(stored_remote_sandbox, runtime)
results.append(result)
return results
def _build_service_url(url: str, service_name: str):
scheme, host_and_path = url.split('://')

View File

@ -755,13 +755,11 @@ class TestSandboxSearch:
sandbox_ids = ['sb1', 'sb2', 'sb3']
mock_response = MagicMock()
mock_response.raise_for_status.return_value = None
mock_response.json.return_value = {
'runtimes': [
create_runtime_data('sb1'),
create_runtime_data('sb2'),
create_runtime_data('sb3'),
]
}
mock_response.json.return_value = [
create_runtime_data('sb1'),
create_runtime_data('sb2'),
create_runtime_data('sb3'),
]
remote_sandbox_service.httpx_client.request = AsyncMock(
return_value=mock_response
)
@ -802,13 +800,11 @@ class TestSandboxSearch:
sandbox_ids = ['sb1', 'sb2', 'sb3']
mock_response = MagicMock()
mock_response.raise_for_status.return_value = None
mock_response.json.return_value = {
'runtimes': [
create_runtime_data('sb1'),
create_runtime_data('sb3'),
# sb2 is missing from the response
]
}
mock_response.json.return_value = [
create_runtime_data('sb1'),
create_runtime_data('sb3'),
# sb2 is missing from the response
]
remote_sandbox_service.httpx_client.request = AsyncMock(
return_value=mock_response
)