compute_engine_base.ts

import { Readable } from 'stream'
import type {
  C2DClusterInfo,
  ComputeEnvironment,
  ComputeAlgorithm,
  ComputeAsset,
  ComputeJob,
  ComputeResourceRequest,
  ComputeResourceRequestWithPrice,
  ComputeResourceType,
  ComputeResource,
  ComputeResourcesPricingInfo,
  DBComputeJobPayment,
  DBComputeJob,
  dockerDeviceRequest,
  DBComputeJobMetadata,
  ComputeEnvFees
} from '../../@types/C2D/C2D.js'
import { C2DClusterType } from '../../@types/C2D/C2D.js'
import { C2DDatabase } from '../database/C2DDatabase.js'
import { Escrow } from '../core/utils/escrow.js'
import { KeyManager } from '../KeyManager/index.js'
import { dockerRegistryAuth, dockerRegistrysAuth } from '../../@types/OceanNode.js'
import { ValidateParams } from '../httpRoutes/validateCommands.js'
import { EncryptMethod } from '../../@types/fileObject.js'
import { CORE_LOGGER } from '../../utils/logging/common.js'
import { DockerRegistryAuthSchema } from '../../utils/config/schemas.js'
export abstract class C2DEngine {
  private clusterConfig: C2DClusterInfo
  public db: C2DDatabase
  public escrow: Escrow
  public keyManager: KeyManager
  public dockerRegistryAuths: dockerRegistrysAuth

  public constructor(
    cluster: C2DClusterInfo,
    db: C2DDatabase,
    escrow: Escrow,
    keyManager: KeyManager,
    dockerRegistryAuths: dockerRegistrysAuth
  ) {
    this.clusterConfig = cluster
    this.db = db
    this.escrow = escrow
    this.keyManager = keyManager
    this.dockerRegistryAuths = dockerRegistryAuths
  }

  getKeyManager(): KeyManager {
    return this.keyManager
  }

  getC2DConfig(): C2DClusterInfo {
    /** Returns cluster config */
    return this.clusterConfig
  }

  getC2DType(): C2DClusterType {
    /** Returns cluster type */
    return this.clusterConfig.type
  }

  // functions which need to be implemented by all engine types
  public abstract getComputeEnvironments(chainId?: number): Promise<ComputeEnvironment[]>

  // overwritten by classes for start actions
  public start(): Promise<void> {
    return null
  }

  // overwritten by classes for cleanup
  public stop(): Promise<void> {
    return null
  }

  // eslint-disable-next-line require-await
  public abstract checkDockerImage(
    image: string,
    encryptedDockerRegistryAuth?: string,
    platform?: any
  ): Promise<ValidateParams>

  public abstract startComputeJob(
    assets: ComputeAsset[],
    algorithm: ComputeAlgorithm,
    output: string,
    environment: string,
    owner: string,
    maxJobDuration: number,
    resources: ComputeResourceRequest[],
    payment: DBComputeJobPayment,
    jobId: string,
    metadata?: DBComputeJobMetadata,
    additionalViewers?: string[],
    queueMaxWaitTime?: number,
    encryptedDockerRegistryAuth?: string
  ): Promise<ComputeJob[]>

  public abstract stopComputeJob(
    jobId: string,
    owner: string,
    agreementId?: string
  ): Promise<ComputeJob[]>

  public abstract getComputeJobStatus(
    consumerAddress?: string,
    agreementId?: string,
    jobId?: string
  ): Promise<ComputeJob[]>

  public abstract getComputeJobResult(
    consumerAddress: string,
    jobId: string,
    index: number,
    offset?: number
  ): Promise<{ stream: Readable; headers: any }>

  public abstract cleanupExpiredStorage(job: DBComputeJob): Promise<boolean>

  public async envExists(
    chainId: number,
    envIdWithHash?: string,
    envIdWithoutHash?: string
  ) {
    try {
      const envs = await this.getComputeEnvironments(chainId)
      for (const c of envs) {
        if (
          (envIdWithHash && c.id === envIdWithHash) ||
          (envIdWithoutHash && this.clusterConfig.hash + '-' + c.id === envIdWithHash)
        ) {
          return true
        }
      }
    } catch (e) {}
    return false
  }

  public async getComputeEnvironment(
    chainId: number,
    envIdWithHash?: string,
    envIdWithoutHash?: string
  ): Promise<ComputeEnvironment> {
    try {
      const envs = await this.getComputeEnvironments(chainId)
      for (const c of envs) {
        if (
          (envIdWithHash && c.id === envIdWithHash) ||
          (envIdWithoutHash && this.clusterConfig.hash + '-' + c.id === envIdWithHash)
        ) {
          return c
        }
      }
    } catch (e) {}
    return null
  }

  public getStreamableLogs(jobId: string): Promise<NodeJS.ReadableStream> {
    throw new Error(`Not implemented for this engine type`)
  }

  protected async getJobEnvironment(job: DBComputeJob): Promise<ComputeEnvironment> {
    const environments: ComputeEnvironment[] = await (
      await this.getComputeEnvironments()
    ).filter((env: ComputeEnvironment) => env.id === job.environment)
    // found it
    if (environments.length === 1) {
      const environment = environments[0]
      return environment
    }
    return null
  }

  /* Returns ComputeResources for a specific resource
   */
  public getMaxMinResource(
    id: ComputeResourceType,
    env: ComputeEnvironment,
    isFree: boolean
  ): ComputeResource {
    const paid = this.getResource(env.resources, id)
    if (!paid) {
      return {
        id,
        total: 0,
        max: 0,
        min: 0
      }
    }
    let free = null
    if (isFree && 'free' in env && 'resources' in env.free) {
      free = this.getResource(env.free.resources, id)
      if (!free) {
        // this resource is not listed under free, so it's not available
        return {
          id,
          total: 0,
          max: 0,
          min: 0
        }
      }
    }
    const total = 'total' in paid ? paid.total : 0
    const max = 'max' in paid ? paid.max : 0
    const min = 'min' in paid ? paid.min : 0
    const ret: ComputeResource = {
      id,
      total: free && 'total' in free ? free.total : total,
      max: free && 'max' in free ? free.max : max,
      min: free && 'min' in free ? free.min : min
    }

    return ret
  }

  // make sure that all requests have cpu, ram, storage
  // eslint-disable-next-line require-await
  public async checkAndFillMissingResources(
    resources: ComputeResourceRequest[],
    env: ComputeEnvironment,
    isFree: boolean
  ): Promise<ComputeResourceRequest[]> {
    if (isFree && !('free' in env)) throw new Error('This env does not support free jobs')
    const properResources: ComputeResourceRequest[] = []
    const elements: string[] = []

    for (const res of isFree ? (env.free?.resources ?? []) : []) elements.push(res.id)
    for (const res of env.resources) if (!elements.includes(res.id)) elements.push(res.id)
    for (const device of elements) {
      let desired = this.getResourceRequest(resources, device)
      const minMax = this.getMaxMinResource(device, env, isFree)
      if (!desired && minMax.min >= 0) {
        // it's required
        desired = minMax.min
      } else {
        if (desired < minMax.min) desired = minMax.min
        if (desired > minMax.max) {
          throw new Error(
            'Not enough ' +
              device +
              ' resources. Requested ' +
              desired +
              ', but max is ' +
              minMax.max
          )
        }
      }
      properResources.push({ id: device, amount: desired })
    }

    this.checkResourceConstraints(properResources, env, isFree)
    return properResources
  }

  protected checkResourceConstraints(
    resources: ComputeResourceRequest[],
    env: ComputeEnvironment,
    isFree: boolean
  ): void {
    const envResources = isFree ? (env.free?.resources ?? []) : (env.resources ?? [])
    for (const envResource of envResources) {
      if (!envResource.constraints || envResource.constraints.length === 0) continue
      const parentAmount = this.getResourceRequest(resources, envResource.id)
      if (!parentAmount || parentAmount <= 0) continue

      for (const constraint of envResource.constraints) {
        let constrainedAmount = this.getResourceRequest(resources, constraint.id) ?? 0

        if (constraint.min !== undefined) {
          const requiredMin = parentAmount * constraint.min
          if (constrainedAmount < requiredMin) {
            const constrainedMaxMin = this.getMaxMinResource(constraint.id, env, isFree)
            if (requiredMin > constrainedMaxMin.max) {
              throw new Error(
                `Cannot satisfy constraint: ${parentAmount} ${envResource.id} requires at least ${requiredMin} ${constraint.id}, but max is ${constrainedMaxMin.max}`
              )
            }
            this.setResourceAmount(resources, constraint.id, requiredMin)
            constrainedAmount = requiredMin
          }
        }

        if (constraint.max !== undefined) {
          const requiredMax = parentAmount * constraint.max
          // re-read in case it was bumped above
          constrainedAmount = this.getResourceRequest(resources, constraint.id) ?? 0
          if (constrainedAmount > requiredMax) {
            throw new Error(
              `Too much ${constraint.id} for ${parentAmount} ${envResource.id}. Max allowed: ${requiredMax}, requested: ${constrainedAmount}`
            )
          }
        }
      }
    }
  }

  protected setResourceAmount(
    resources: ComputeResourceRequest[],
    id: ComputeResourceType,
    amount: number
  ): void {
    for (const resource of resources) {
      if (resource.id === id) {
        resource.amount = amount
        return
      }
    }
  }

  public async getUsedResources(env: ComputeEnvironment): Promise<any> {
    const usedResources: { [x: string]: any } = {}
    const usedFreeResources: { [x: string]: any } = {}
    let jobs: DBComputeJob[] = []
    try {
      jobs = await this.db.getRunningJobs(this.getC2DConfig().hash)
    } catch (e) {
      CORE_LOGGER.error('Failed to get running jobs:' + e.message)
    }

    const envResourceMap = new Map((env.resources || []).map((r) => [r.id, r]))

    let totalJobs = 0
    let totalFreeJobs = 0
    let queuedJobs = 0
    let queuedFreeJobs = 0
    let maxWaitTime = 0
    let maxWaitTimeFree = 0
    let maxRunningTime = 0
    let maxRunningTimeFree = 0

    for (const job of jobs) {
      const isThisEnv = job.environment === env.id
      const isRunning = job.queueMaxWaitTime === 0

      if (isThisEnv) {
        if (isRunning) {
          const timeElapsed = job.buildStartTimestamp
            ? new Date().getTime() / 1000 - Number.parseFloat(job?.buildStartTimestamp)
            : new Date().getTime() / 1000 - Number.parseFloat(job?.algoStartTimestamp)
          totalJobs++
          maxRunningTime += job.maxJobDuration - timeElapsed
          if (job.isFree) {
            totalFreeJobs++
            maxRunningTimeFree += job.maxJobDuration - timeElapsed
          }
        } else {
          queuedJobs++
          maxWaitTime += job.maxJobDuration
          if (job.isFree) {
            queuedFreeJobs++
            maxWaitTimeFree += job.maxJobDuration
          }
        }
      }

      if (isRunning) {
        for (const resource of job.resources) {
          const envRes = envResourceMap.get(resource.id)
          if (envRes) {
            // GPUs are shared-exclusive: inUse tracked globally across all envs
            // Everything else (cpu, ram, disk) is per-env exclusive
            const isSharedExclusive = envRes.type === 'gpu'
            if (!isSharedExclusive && !isThisEnv) continue
            if (!(resource.id in usedResources)) usedResources[resource.id] = 0
            usedResources[resource.id] += resource.amount
            if (job.isFree) {
              if (!(resource.id in usedFreeResources)) usedFreeResources[resource.id] = 0
              usedFreeResources[resource.id] += resource.amount
            }
          }
        }
      }
    }
    return {
      totalJobs,
      totalFreeJobs,
      usedResources,
      usedFreeResources,
      queuedJobs,
      queuedFreeJobs,
      maxWaitTime,
      maxWaitTimeFree,
      maxRunningTime,
      maxRunningTimeFree
    }
  }

  protected physicalLimits: Map<string, number> = new Map()

  private checkGlobalResourceAvailability(
    allEnvironments: ComputeEnvironment[],
    resourceId: string,
    amount: number
  ) {
    let globalUsed = 0
    let globalTotal = 0
    for (const e of allEnvironments) {
      const res = this.getResource(e.resources, resourceId)
      if (res) {
        globalTotal += res.total || 0
        globalUsed += res.inUse || 0
      }
    }
    const physicalLimit = this.physicalLimits.get(resourceId)
    if (physicalLimit !== undefined && globalTotal > physicalLimit) {
      globalTotal = physicalLimit
    }
    const globalRemainder = globalTotal - globalUsed
    if (globalRemainder < amount) {
      throw new Error(
        `Not enough available ${resourceId} globally (remaining: ${globalRemainder}, requested: ${amount})`
      )
    }
  }

  // overridden by each engine if required
  // eslint-disable-next-line require-await
  public async checkIfResourcesAreAvailable(
    resourcesRequest: ComputeResourceRequest[],
    env: ComputeEnvironment,
    isFree: boolean,
    allEnvironments?: ComputeEnvironment[]
  ) {
    // Filter out resources with amount 0 as they're not actually being requested
    const activeResources = resourcesRequest.filter((r) => r.amount > 0)

    for (const request of activeResources) {
      let envResource = this.getResource(env.resources, request.id)
      if (!envResource) throw new Error(`No such resource ${request.id}`)
      if (envResource.total - envResource.inUse < request.amount)
        throw new Error(`Not enough available ${request.id}`)

      // Global check for non-GPU resources (cpu, ram, disk are per-env exclusive)
      // GPUs are shared-exclusive so their inUse already reflects global usage
      if (allEnvironments && envResource.type !== 'gpu') {
        this.checkGlobalResourceAvailability(allEnvironments, request.id, request.amount)
      }

      if (isFree) {
        if (!env.free) throw new Error(`No free resources`)
        envResource = this.getResource(env.free?.resources, request.id)
        if (!envResource) throw new Error(`No such free resource ${request.id}`)
        if (envResource.total - envResource.inUse < request.amount)
          throw new Error(`Not enough available ${request.id} for free`)
      }
    }
    if ('maxJobs' in env && env.maxJobs && env.runningJobs + 1 > env.maxJobs) {
      throw new Error(`Too many running jobs `)
    }
    if (
      isFree &&
      'free' in env &&
      `maxJobs` in env.free &&
      env.free.maxJobs &&
      env.runningfreeJobs + 1 > env.free.maxJobs
    ) {
      throw new Error(`Too many running free jobs `)
    }
  }

  public getResource(resources: ComputeResource[], id: ComputeResourceType) {
    if (!resources) return null
    for (const resource of resources) {
      if (resource.id === id) {
        return resource
      }
    }
    return null
  }

  public getResourceRequest(
    resources: ComputeResourceRequest[],
    id: ComputeResourceType
  ) {
    if (!resources) return null
    for (const resource of resources) {
      if (resource.id === id) {
        return resource.amount
      }
    }
    return null
  }

  public getDockerDeviceRequest(
    requests: ComputeResourceRequest[],
    resources: ComputeResource[]
  ): dockerDeviceRequest[] | null {
    if (!resources) return null

    // Filter out resources with amount 0 as they're not actually being requested
    const activeResources = requests.filter((r) => r.amount > 0)
    const grouped: Record<string, dockerDeviceRequest> = {}

    for (const resource of activeResources) {
      const res = this.getResource(resources, resource.id)
      const init = res?.init?.deviceRequests
      if (!init) continue

      const key = `${init.Driver}-${JSON.stringify(init.Capabilities)}`
      if (!grouped[key]) {
        grouped[key] = {
          Driver: init.Driver,
          Capabilities: init.Capabilities,
          DeviceIDs: [],
          Options: init.Options ?? null,
          Count: undefined
        }
      }

      if (init.DeviceIDs?.length) {
        grouped[key].DeviceIDs!.push(...init.DeviceIDs)
      }
    }

    return Object.values(grouped)
  }

  public getDockerAdvancedConfig(
    requests: ComputeResourceRequest[],
    resources: ComputeResource[]
  ) {
    const ret = {
      Devices: [] as any[],
      GroupAdd: [] as string[],
      SecurityOpt: [] as string[],
      Binds: [] as string[],
      CapAdd: [] as string[],
      CapDrop: [] as string[],
      IpcMode: null as string,
      ShmSize: 0 as number
    }
    // Filter out resources with amount 0 as they're not actually being requested
    const activeResources = requests.filter((r) => r.amount > 0)

    for (const resource of activeResources) {
      const res = this.getResource(resources, resource.id)
      if (res.init && res.init.advanced) {
        for (const [key, value] of Object.entries(res.init.advanced)) {
          switch (key) {
            case 'IpcMode':
              ret.IpcMode = value as string
              break
            case 'ShmSize':
              ret.ShmSize = value as number
              break
            case 'GroupAdd':
              for (const grp of value as string[]) {
                if (!ret.GroupAdd.includes(grp)) ret.GroupAdd.push(grp)
              }
              break
            case 'CapAdd':
              for (const grp of value as string[]) {
                if (!ret.CapAdd.includes(grp)) ret.CapAdd.push(grp)
              }
              break
            case 'CapDrop':
              for (const grp of value as string[]) {
                if (!ret.CapDrop.includes(grp)) ret.CapDrop.push(grp)
              }
              break
            case 'Devices':
              for (const device of value as string[]) {
                if (!ret.Devices.find((d) => d.PathOnHost === device))
                  ret.Devices.push({
                    PathOnHost: device,
                    PathInContainer: device,
                    CgroupPermissions: 'rwm'
                  })
              }
              break
            case 'SecurityOpt':
              for (const [secKeys, secValues] of Object.entries(value))
                if (!ret.SecurityOpt.includes(secKeys + '=' + secValues))
                  ret.SecurityOpt.push(secKeys + '=' + secValues)
              break
            case 'Binds':
              for (const grp of value as string[]) {
                if (!ret.Binds.includes(grp)) ret.Binds.push(grp)
              }
              break
          }
        }
      }
    }
    return ret
  }

  public getEnvPricesForToken(
    env: ComputeEnvironment,
    chainId: number,
    token: string
  ): ComputeResourcesPricingInfo[] {
    if (!env.fees || !(chainId in env.fees) || !env.fees[chainId]) {
      return null
    }
    for (const fee of env.fees[chainId]) {
      // eslint-disable-next-line security/detect-possible-timing-attacks
      if (fee.feeToken === token) {
        return fee.prices
      }
    }

    return null
  }

  public getResourcePrice(
    prices: ComputeResourcesPricingInfo[],
    id: ComputeResourceType
  ) {
    for (const pr of prices) {
      if (pr.id === id) {
        return pr.price
      }
    }
    return 0
  }

  public getTotalCostOfJob(
    resources: ComputeResourceRequestWithPrice[],
    duration: number,
    fee: ComputeEnvFees
  ) {
    let cost: number = 0
    for (const request of resources) {
      const price = fee.prices.find((p) => p.id === request.id)?.price
      if (price) {
        cost += price * request.amount * Math.ceil(duration / 60)
      }
    }
    return cost
  }

  public calculateResourcesCost(
    resourcesRequest: ComputeResourceRequest[],
    env: ComputeEnvironment,
    chainId: number,
    token: string,
    maxJobDuration: number
  ): number | null {
    if (maxJobDuration < env.minJobDuration) maxJobDuration = env.minJobDuration
    const prices = this.getEnvPricesForToken(env, chainId, token)
    if (!prices) return null
    let cost: number = 0
    for (const request of resourcesRequest) {
      const resourcePrice = this.getResourcePrice(prices, request.id)
      cost += resourcePrice * request.amount * Math.ceil(maxJobDuration / 60)
    }
    return cost
  }

  public getDockerRegistryAuth(registry: string): dockerRegistryAuth | null {
    if (!this.dockerRegistryAuths) return null
    if (this.dockerRegistryAuths[registry]) {
      return this.dockerRegistryAuths[registry]
    }
    return null
  }

  public async checkEncryptedDockerRegistryAuth(
    encryptedDockerRegistryAuth: string
  ): Promise<ValidateParams> {
    let decryptedDockerRegistryAuth: dockerRegistryAuth
    try {
      const decryptedDockerRegistryAuthBuffer = await this.keyManager.decrypt(
        Uint8Array.from(Buffer.from(encryptedDockerRegistryAuth, 'hex')),
        EncryptMethod.ECIES
      )

      // Convert decrypted buffer to string and parse as JSON
      const decryptedDockerRegistryAuthString =
        decryptedDockerRegistryAuthBuffer.toString()

      decryptedDockerRegistryAuth = JSON.parse(decryptedDockerRegistryAuthString)
    } catch (error: any) {
      const errorMessage = `Invalid encryptedDockerRegistryAuth: failed to parse JSON - ${error?.message || String(error)}`
      CORE_LOGGER.error(errorMessage)
      return {
        valid: false,
        reason: errorMessage,
        status: 400
      }
    }

    // Validate using schema - ensures either auth or username+password are provided
    const validationResult = DockerRegistryAuthSchema.safeParse(
      decryptedDockerRegistryAuth
    )
    if (!validationResult.success) {
      const errorMessageValidation = validationResult.error.errors
        .map((err) => err.message)
        .join('; ')
      const errorMessage = `Invalid encryptedDockerRegistryAuth: ${errorMessageValidation}`
      CORE_LOGGER.error(errorMessage)
      return {
        valid: false,
        reason: errorMessage,
        status: 400
      }
    }
    return {
      valid: true,
      reason: null,
      status: 200
    }
  }
}