fix: harden registry model sync timeouts (#2407)

* fix: harden registry model sync timeouts

* fix: reconnect after ambiguous registry add

Author: yuranich

.github/workflows/pr-models-validation-registry-server.yml

@@ -256,7 +256,7 @@ jobs:
         working-directory: ${{ env.PKG_DIR }}
         shell: bash
         run: node scripts/sync-models.js --file=./data/models.prod.json
-        timeout-minutes: 30
+        timeout-minutes: 70
 
   smoke-test:
     needs:

packages/registry-server/lib/registry-service.js

@@ -210,6 +210,13 @@ class RegistryService extends ReadyResource {
       const peerKey = peerInfo?.publicKey ? IdEnc.normalize(peerInfo.publicKey) : null
 
       this.logger.info({ peer: peerKey || 'unknown' }, 'Swarm connection opened')
+      conn.on('error', err => {
+        this.logger.warn({
+          peer: peerKey || 'unknown',
+          error: err.message,
+          code: err.code
+        }, 'Swarm connection error')
+      })
       conn.on('close', () => {
         this.logger.info({ peer: peerKey || 'unknown' }, 'Swarm connection closed')
       })

packages/registry-server/scripts/sync-models.js

@@ -9,6 +9,9 @@ const { connectToRegistry } = require('./utils/rpc-client')
 const { parseCanonicalSource } = require('../lib/source-helpers')
 const QVACRegistryClient = require('../client/lib/client')
 
+const ADD_MODEL_RPC_TIMEOUT_MS = 60 * 60 * 1000
+const ADD_MODEL_POLL_INTERVAL_MS = 10 * 1000
+
 async function syncModels () {
   const args = process.argv.slice(2)
   const fileArg = args.find(arg => arg.startsWith('--file='))
@@ -30,7 +33,7 @@ async function syncModels () {
   const client = new QVACRegistryClient({ registryCoreKey, logger })
   await client.ready()
 
-  const connection = await connectToRegistry({ config, logger })
+  let connection = await connectToRegistry({ config, logger })
 
   try {
     const configModels = JSON.parse(await fs.readFile(path.resolve(filePath), 'utf8'))
@@ -95,7 +98,30 @@ async function syncModels () {
               modelRequest.deprecationReason = entry.deprecationReason
             }
 
-            await connection.rpc.request('add-model', modelRequest)
+            try {
+              await connection.rpc.request('add-model', modelRequest, { timeout: ADD_MODEL_RPC_TIMEOUT_MS })
+            } catch (err) {
+              if (!isAmbiguousRpcError(err)) throw err
+
+              logger.warn({
+                path: sourceInfo.path,
+                source: sourceInfo.protocol,
+                error: err.message,
+                code: err.code
+              }, 'add-model RPC ended ambiguously; polling registry for completed ingest')
+
+              const recovery = await recoverAfterAmbiguousAdd({
+                client,
+                sourceInfo,
+                logger,
+                connection,
+                reconnect: () => connectToRegistry({ config, logger })
+              })
+              connection = recovery.connection
+
+              if (recovery.error) throw recovery.error
+              dbByKey.set(key, recovery.model)
+            }
           } else {
             logger.info(`[DRY RUN] Would add: ${sourceInfo.path}`)
           }
@@ -185,6 +211,88 @@ async function syncModels () {
   }
 }
 
+async function recoverAfterAmbiguousAdd ({
+  client,
+  sourceInfo,
+  logger,
+  connection,
+  reconnect,
+  waitForModel = waitForModelAfterAmbiguousAdd
+}) {
+  let model = null
+  let error = null
+
+  try {
+    model = await waitForModel({
+      client,
+      sourceInfo,
+      logger
+    })
+  } catch (err) {
+    error = err
+  }
+
+  await connection.cleanup().catch(cleanupErr => {
+    logger.warn({ error: cleanupErr.message }, 'Failed to clean up stale RPC connection')
+  })
+
+  return {
+    connection: await reconnect(),
+    error,
+    model
+  }
+}
+
+function isAmbiguousRpcError (err) {
+  if (!err) return false
+
+  const code = err.code || err.cause?.code
+  if (code === 'ETIMEDOUT' || code === 'CHANNEL_CLOSED' || code === 'CHANNEL_DESTROYED') {
+    return true
+  }
+
+  const message = String(err.message || '').toLowerCase()
+  return (
+    message.includes('connection timed out') ||
+    message.includes('channel closed') ||
+    message.includes('channel destroyed')
+  )
+}
+
+async function waitForModelAfterAmbiguousAdd ({
+  client,
+  sourceInfo,
+  timeoutMs = ADD_MODEL_RPC_TIMEOUT_MS,
+  pollIntervalMs = ADD_MODEL_POLL_INTERVAL_MS,
+  logger = console,
+  sleep = defaultSleep
+}) {
+  const startedAt = Date.now()
+
+  while (Date.now() - startedAt <= timeoutMs) {
+    if (client.db?.core?.update) {
+      await client.db.core.update()
+    }
+
+    const model = await client.getModel(sourceInfo.path, sourceInfo.protocol)
+    if (model) {
+      logger.info({
+        path: sourceInfo.path,
+        source: sourceInfo.protocol
+      }, 'Model appeared after ambiguous add-model RPC')
+      return model
+    }
+
+    await sleep(pollIntervalMs)
+  }
+
+  throw new Error(`Timed out waiting for model after ambiguous add-model RPC: ${sourceInfo.path}`)
+}
+
+function defaultSleep (ms) {
+  return new Promise(resolve => setTimeout(resolve, ms))
+}
+
 function needsMetadataUpdate (config, existing, sourceInfo) {
   return (
     config.engine !== existing.engine ||
@@ -274,4 +382,10 @@ if (require.main === module) {
     })
 }
 
-module.exports = { syncModels }
+module.exports = {
+  ADD_MODEL_RPC_TIMEOUT_MS,
+  recoverAfterAmbiguousAdd,
+  isAmbiguousRpcError,
+  syncModels,
+  waitForModelAfterAmbiguousAdd
+}

packages/registry-server/scripts/utils/rpc-client.js

@@ -97,6 +97,14 @@ async function connectToRegistry ({ config, logger = console, storage = './temp-
 
       logger.info('RPC Client: Connected to server', { peer: peerKey })
 
+      conn.on('error', err => {
+        logger.warn({
+          peer: peerKey,
+          error: err.message,
+          code: err.code
+        }, 'RPC Client: connection error')
+      })
+
       const rpc = new ProtomuxRPC(conn, {
         protocol: 'qvac-registry-rpc',
         valueEncoding: cenc.json

packages/registry-server/tests/unit/sync-models.test.js

@@ -0,0 +1,80 @@
+'use strict'
+
+const test = require('brittle')
+const {
+  ADD_MODEL_RPC_TIMEOUT_MS,
+  recoverAfterAmbiguousAdd,
+  isAmbiguousRpcError,
+  waitForModelAfterAmbiguousAdd
+} = require('../../scripts/sync-models')
+
+test('add-model RPC timeout is one hour', t => {
+  t.is(ADD_MODEL_RPC_TIMEOUT_MS, 60 * 60 * 1000)
+})
+
+test('isAmbiguousRpcError identifies transport timeouts and channel closes', t => {
+  t.ok(isAmbiguousRpcError(Object.assign(new Error('connection timed out'), { code: 'ETIMEDOUT' })))
+  t.ok(isAmbiguousRpcError(Object.assign(new Error('closed'), { code: 'CHANNEL_CLOSED' })))
+  t.ok(isAmbiguousRpcError(new Error('Channel closed')))
+  t.absent(isAmbiguousRpcError(new Error('License not found')))
+})
+
+test('waitForModelAfterAmbiguousAdd polls until the model appears', async t => {
+  const expected = { path: 'repo/model.gguf', source: 'hf' }
+  const calls = []
+
+  const client = {
+    async getModel (modelPath, source) {
+      calls.push([modelPath, source])
+      return calls.length === 2 ? expected : null
+    }
+  }
+
+  const result = await waitForModelAfterAmbiguousAdd({
+    client,
+    sourceInfo: { path: expected.path, protocol: expected.source },
+    timeoutMs: 10,
+    pollIntervalMs: 5,
+    logger: { info () {} },
+    sleep: async () => {}
+  })
+
+  t.alike(result, expected)
+  t.alike(calls, [
+    [expected.path, expected.source],
+    [expected.path, expected.source]
+  ])
+})
+
+test('recoverAfterAmbiguousAdd reconnects even when polling times out', async t => {
+  t.plan(4)
+
+  const staleConnection = {
+    cleaned: false,
+    async cleanup () {
+      this.cleaned = true
+    }
+  }
+  const freshConnection = {}
+  const pollError = new Error('poll timed out')
+  let reconnects = 0
+
+  const result = await recoverAfterAmbiguousAdd({
+    client: {},
+    sourceInfo: { path: 'repo/model.gguf', protocol: 'hf' },
+    logger: { info () {}, warn () {} },
+    connection: staleConnection,
+    reconnect: async () => {
+      reconnects++
+      return freshConnection
+    },
+    waitForModel: async () => {
+      throw pollError
+    }
+  })
+
+  t.is(result.error, pollError)
+  t.is(result.connection, freshConnection)
+  t.is(reconnects, 1)
+  t.ok(staleConnection.cleaned)
+})