feat(cu): make llama module apply configurable via MODULE_MODE

servers/cu/src/config.js

@@ -91,6 +91,7 @@ const preprocessedServerConfigSchema = z.preprocess(
 const CONFIG_ENVS = {
   development: {
     MODE,
+    MODULE_MODE: process.env.MODULE_MODE,
     port: process.env.PORT || 6363,
     ENABLE_METRICS_ENDPOINT: process.env.ENABLE_METRICS_ENDPOINT,
     GATEWAY_URL: process.env.GATEWAY_URL || 'https://arweave.net',
@@ -138,6 +139,7 @@ const CONFIG_ENVS = {
   },
   production: {
     MODE,
+    MODULE_MODE: process.env.MODULE_MODE,
     port: process.env.PORT || 6363,
     ENABLE_METRICS_ENDPOINT: process.env.ENABLE_METRICS_ENDPOINT,
     GATEWAY_URL: process.env.GATEWAY_URL || 'https://arweave.net',

servers/cu/src/domain/index.js

@@ -318,14 +318,7 @@ export const createApis = async (ctx) => {
     isModuleComputeLimitSupported: WasmClient.isModuleComputeLimitSupportedWith({ PROCESS_WASM_COMPUTE_MAX_LIMIT: ctx.PROCESS_WASM_COMPUTE_MAX_LIMIT }),
     isModuleFormatSupported: WasmClient.isModuleFormatSupportedWith({ PROCESS_WASM_SUPPORTED_FORMATS: ctx.PROCESS_WASM_SUPPORTED_FORMATS }),
     isModuleExtensionSupported: WasmClient.isModuleExtensionSupportedWith({ PROCESS_WASM_SUPPORTED_EXTENSIONS: ctx.PROCESS_WASM_SUPPORTED_EXTENSIONS }),
-    /**
-     * This is not configurable, as Load message support will be dictated by the protocol,
-     * which proposes Assignments as a more flexible and extensible solution that also includes
-     * Load Message use-cases
-     *
-     * But for now, supporting Load messages in perpetuity.
-     */
-    AO_LOAD_MAX_BLOCK: Number.MAX_SAFE_INTEGER,
+    MODULE_MODE: ctx.MODULE_MODE,
     logger
   })
   /**

servers/cu/src/domain/lib/loadModule.js

@@ -7,6 +7,17 @@ import { findModuleSchema, loadTransactionMetaSchema, saveModuleSchema } from '.
 import { addressFrom, parseTags } from '../utils.js'
 import { rawTagSchema } from '../model.js'
 
+export const LLAMA_ADMISSABLE = [
+  'dx3GrOQPV5Mwc1c-4HTsyq0s1TNugMf7XfIKJkyVQt8', // Random NFT metadata (1.7kb of JSON)
+  'XOJ8FBxa6sGLwChnxhF2L71WkKLSKq1aU5Yn5WnFLrY', // GPT-2 117M model.
+  'M-OzkyjxWhSvWYF87p0kvmkuAEEkvOzIj4nMNoSIydc', // GPT-2-XL 4-bit quantized model.
+  'kd34P4974oqZf2Db-hFTUiCipsU6CzbR6t-iJoQhKIo', // Phi-2
+  'ISrbGzQot05rs_HKC08O_SmkipYQnqgB1yC3mjZZeEo', // Phi-3 Mini 4k Instruct
+  'sKqjvBbhqKvgzZT4ojP1FNvt4r_30cqjuIIQIr-3088', // CodeQwen 1.5 7B Chat q3
+  'Pr2YVrxd7VwNdg6ekC0NXWNKXxJbfTlHhhlrKbAd1dA', // Llama3 8B Instruct q4
+  'jbx-H6aq7b3BbNCHlK50Jz9L-6pz9qmldrYXMwjqQVI' // Llama3 8B Instruct q8
+]
+
 /**
  * The result that is produced from this step
  * and added to ctx.
@@ -88,7 +99,7 @@ function getModuleWith ({ findModule, saveModule, loadTransactionMeta, logger })
   }
 }
 
-function setModuleOptionsWith ({ isModuleMemoryLimitSupported, isModuleComputeLimitSupported, isModuleFormatSupported, isModuleExtensionSupported }) {
+function setModuleOptionsWith ({ isModuleMemoryLimitSupported, isModuleComputeLimitSupported, isModuleFormatSupported, isModuleExtensionSupported, MODULE_MODE }) {
   const checkModuleOption = (name, pred, err) => (options) =>
     of()
       .chain(fromPromise(async () => pred(options[name])))
@@ -158,11 +169,18 @@ function setModuleOptionsWith ({ isModuleMemoryLimitSupported, isModuleComputeLi
           {}
         )
       ),
-      mode: (args) => pathOr(
-        pathOr('Assignments', ['moduleTags', 'Availability-Type'], args),
-        ['processTags', 'Availability-Type'],
-        args
-      ),
+      mode: (args) => {
+        if (MODULE_MODE === 'llama') return 'test'
+
+        return pathOr(
+          pathOr('Assignments', ['moduleTags', 'Availability-Type'], args),
+          ['processTags', 'Availability-Type'],
+          args
+        )
+      },
+      admissableList: () => {
+        if (MODULE_MODE === 'llama') return LLAMA_ADMISSABLE
+      },
       spawn: () => ({ id: ctx.id, owner: ctx.owner, tags: ctx.tags }),
       module: () => ({ id: ctx.moduleId, owner: ctx.moduleOwner, tags: ctx.moduleTags }),
       /**

servers/cu/src/domain/lib/loadModule.test.js

@@ -5,7 +5,7 @@ import * as assert from 'node:assert'
 import { lensIndex, remove, set } from 'ramda'
 
 import { createLogger } from '../logger.js'
-import { loadModuleWith } from './loadModule.js'
+import { LLAMA_ADMISSABLE, loadModuleWith } from './loadModule.js'
 
 const PROCESS = 'contract-123-9HdeqeuYQOgMgWucro'
 const logger = createLogger('ao-cu:readState')
@@ -79,6 +79,62 @@ describe('loadModule', () => {
     assert.equal(result.id, PROCESS)
   })
 
+  test('use llama mode and admissable list if the MODULE_MODE is llama', async () => {
+    const loadModule = loadModuleWith({
+      loadTransactionMeta: async () => ({
+        owner: moduleOwner,
+        tags: moduleTags
+      }),
+      findModule: async () => { throw { status: 404 } },
+      saveModule: async () => 'foobar',
+      isModuleMemoryLimitSupported: async ({ limit }) => {
+        assert.equal(limit, 11264)
+        return true
+      },
+      isModuleComputeLimitSupported: async ({ limit }) => {
+        assert.equal(limit, 10000)
+        return true
+      },
+      isModuleFormatSupported: async ({ format }) => {
+        assert.equal(format, 'wasm32-unknown-emscripten')
+        return true
+      },
+      isModuleExtensionSupported: async ({ extension }) => true,
+      MODULE_MODE: 'llama',
+      logger
+    })
+
+    const processTags = [
+      { name: 'Module', value: 'foobar' },
+      { name: 'Compute-Limit', value: '10000' },
+      { name: 'Memory-Limit', value: '11-kb' }
+    ]
+
+    const result = await loadModule({
+      id: PROCESS,
+      tags: processTags,
+      owner: 'p-owner-123'
+    }).toPromise()
+
+    assert.equal(result.moduleId, 'foobar')
+    assert.deepStrictEqual(result.moduleTags, moduleTags)
+    assert.equal(result.moduleOwner, 'owner-123')
+    assert.deepStrictEqual(result.moduleOptions, {
+      format: 'wasm32-unknown-emscripten',
+      inputEncoding: 'JSON-1',
+      outputEncoding: 'JSON-1',
+      computeLimit: 10000,
+      memoryLimit: 11264,
+      extensions: {},
+      mode: 'test',
+      admissableList: LLAMA_ADMISSABLE,
+      spawn: { id: PROCESS, owner: 'p-owner-123', tags: processTags },
+      module: { id: 'foobar', owner: 'owner-123', tags: moduleTags },
+      blockHeight: 100
+    })
+    assert.equal(result.id, PROCESS)
+  })
+
   test('should fallback to Module tags for Compute-Limit and Module-Limit if not on Process', async () => {
     const loadModule = loadModuleWith({
       loadTransactionMeta: async () => ({

servers/cu/src/domain/model.js

@@ -14,6 +14,10 @@ export const commaDelimitedArraySchema = z.preprocess((val) => {
 }, z.array(z.string()))
 
 export const domainConfigSchema = z.object({
+  /**
+   * shim to support llama configurability.
+   */
+  MODULE_MODE: z.string().nullish(),
   /**
    * The maximum Memory-Limit, in bytes, supported for ao processes
    *