1204 break import cycles

.github/workflows/_build.yml

@@ -3,37 +3,6 @@ on:
   workflow_call:
 
 jobs:
-  build-lib:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F discojs run build
-
-  build-lib-node:
-    needs: build-lib
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F discojs-node run build
-
-  build-lib-web:
-    needs: build-lib
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F discojs-web run build
-
-  build-server:
-    needs: build-lib-node
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F server run build
-
   build-server-docker:
     runs-on: ubuntu-latest
     steps:
@@ -51,29 +20,11 @@ jobs:
             timeout=$((timeout - 1))
           done
 
-  build-cli:
-    needs: build-server
+  build-topological:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v7
       - uses: ./.github/actions/setup-node
-      # Build all subprojects to upload all artifacts at once
       - run: pnpm -r run build
       - uses: actions/upload-artifact@v7
         with: { name: all-builds, path: "*/dist" }
-
-  build-webapp:
-    needs: build-lib-web
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F webapp run build
-
-  build-docs-examples:
-    needs: build-server
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v7
-      - uses: ./.github/actions/setup-node
-      - run: pnpm -F examples run build

.github/workflows/_static-analysis.yml

@@ -16,3 +16,10 @@ jobs:
       - uses: actions/checkout@v7
       - uses: ./.github/actions/setup-node
       - run: pnpm exec knip
+
+  check_cycles:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: ./.github/actions/setup-node
+      - run: pnpm run check_cycles

.gitignore

@@ -6,6 +6,7 @@
 # stored trained models
 /server/models/
 /models/
+/docs/examples/models/
 
 # tsc built
 /*/dist/

.knip.json

@@ -5,7 +5,6 @@
       "ignoreDependencies": ["typescript-eslint", "supports-color"]
     },
     "webapp": {
-      "ignore": ["env.d.ts"],
       "ignoreDependencies": ["buffer"]
     },
     "isomorphic-wrtc": {},

Dockerfile

@@ -18,6 +18,7 @@ RUN pnpm ci
 COPY discojs/ discojs/
 COPY discojs-node/ discojs-node/
 COPY tsconfig.base.json .
+COPY tsconfig.base.lib.json .
 RUN pnpm -F discojs -F discojs-node run build
 
 COPY server/ server/

cli/src/args.ts

@@ -279,6 +279,6 @@ export const args: BenchmarkArguments = {
 
       return task;
     },
-    getModel: () => provider.getModel(),
+    modelCard: provider.modelCard,
   },
 };

cli/src/benchmark_gpt.ts

@@ -2,12 +2,13 @@ import "@tensorflow/tfjs-node";
 import { List } from "immutable";
 import { parse } from "ts-command-line-args";
 
-import type { Network, Task } from "@epfml/discojs";
+import type { Network, Task, GPTConfig } from "@epfml/discojs";
 import {
-  async_iterator,
+  gather,
   defaultTasks,
+  defaultModels,
   fetchTasks,
-  models,
+  GPT,
 } from "@epfml/discojs";
 import { loadModelFromDisk, loadText } from "@epfml/discojs-node";
 
@@ -86,7 +87,10 @@ async function main(args: Required<CLIArguments>): Promise<void> {
   } = args;
 
   // Launch a server instance
-  const server = await Server.with(defaultTasks.wikitext);
+  const server = await Server.with(
+    [defaultModels.Wikitext],
+    [defaultTasks.wikitext],
+  );
   const [handle, url] = await server.serve();
 
   // Fetch the wikitext task from the server
@@ -105,8 +109,8 @@ async function main(args: Required<CLIArguments>): Promise<void> {
     const epochsCount = 1;
     const iterationsPerEpoch = 10;
 
-    const config: models.GPTConfig = {
-      modelType: modelType as models.GPTConfig["modelType"],
+    const config: GPTConfig = {
+      modelType: modelType as GPTConfig["modelType"],
       maxIter: iterationsPerEpoch,
       lr: 0.0001,
       contextLength,
@@ -126,16 +130,14 @@ async function main(args: Required<CLIArguments>): Promise<void> {
       .batch(batchSize);
 
     // Init and train the model
-    const model = new models.GPT(config);
+    const model = new GPT(config);
     console.log(
       `\tmodel type ${modelType} \n\tbatch size ${batchSize} \n\tcontext length ${contextLength}`,
     );
 
     let epochTime = performance.now();
     for (let epochsCounter = 1; epochsCounter <= epochsCount; epochsCounter++) {
-      const [_, logs] = await async_iterator.gather(
-        model.train(preprocessedDataset),
-      );
+      const [_, logs] = await gather(model.train(preprocessedDataset));
       epochTime = performance.now() - epochTime;
       const msPerToken =
         epochTime /
@@ -150,7 +152,7 @@ async function main(args: Required<CLIArguments>): Promise<void> {
      */
   } else {
     const model = await loadModelFromDisk(modelPath);
-    if (!(model instanceof models.GPT)) {
+    if (!(model instanceof GPT)) {
       throw new Error("Loaded model isn't a GPT model");
     }
 

cli/src/cli.ts

@@ -15,11 +15,7 @@ import type {
   TaskProvider,
   Network,
 } from "@epfml/discojs";
-import {
-  Disco,
-  aggregator as aggregators,
-  client as clients,
-} from "@epfml/discojs";
+import { Disco, getAggregator, getClient } from "@epfml/discojs";
 
 import { getTaskData } from "./data.js";
 import { args } from "./args.js";
@@ -35,8 +31,8 @@ async function runUser<D extends DataType, N extends Network>(
 ): Promise<List<SummaryLogs>> {
   // cast as typescript isn't good with generics
   const trainingScheme = task.trainingInformation.scheme as N;
-  const aggregator = aggregators.getAggregator(task);
-  const client = clients.getClient(trainingScheme, url, task, aggregator);
+  const aggregator = getAggregator(task);
+  const client = getClient(trainingScheme, url, task, aggregator);
   const disco = new Disco(task, client, { scheme: trainingScheme });
 
   const dir = path.join(".", `${args.testID}`);

cli/src/data.ts

@@ -1,6 +1,7 @@
 import path from "node:path";
-import { Dataset, processing } from "@epfml/discojs";
-import { DataFormat, DataType, Image, Task } from "@epfml/discojs";
+import type { Dataset } from "@epfml/discojs";
+import { extractColumn } from "@epfml/discojs";
+import type { DataFormat, DataType, Image, Task } from "@epfml/discojs";
 import { loadCSV, loadImage, loadImagesInDir } from "@epfml/discojs-node";
 import { Repeat } from "immutable";
 
@@ -49,10 +50,7 @@ function loadTinderDogData(split: number): Dataset<DataFormat.Raw["image"]> {
   return loadCSV(path.join(folder, "labels.csv"))
     .map(
       (row) =>
-        [
-          processing.extractColumn(row, "filename"),
-          processing.extractColumn(row, "label"),
-        ] as const,
+        [extractColumn(row, "filename"), extractColumn(row, "label")] as const,
     )
     .map(async ([filename, label]) => {
       try {
@@ -76,10 +74,7 @@ function loadData(
   return loadCSV(path.join(folder, "labels.csv"))
     .map(
       (row) =>
-        [
-          processing.extractColumn(row, "filename"),
-          processing.extractColumn(row, "label"),
-        ] as const,
+        [extractColumn(row, "filename"), extractColumn(row, "label")] as const,
     )
     .map(async ([filename, label]) => {
       try {

cli/src/hellaswag_gpt.ts

@@ -5,22 +5,25 @@ import { parse } from "ts-command-line-args";
 
 import "@tensorflow/tfjs-node";
 import path from "node:path";
-import { models, serialization, Tokenizer } from "@epfml/discojs";
+import type { HellaSwagDataset } from "@epfml/discojs";
+import {
+  GPT,
+  ONNXModel,
+  modelDecode,
+  Tokenizer,
+  evaluate_hellaswag,
+} from "@epfml/discojs";
 import { loadHellaSwag } from "@epfml/discojs-node";
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
 
-async function evaluateModel(
-  model: models.GPT | models.ONNXModel,
-  numDataPoints = -1,
-) {
-  const hellaswagDataset: models.HellaSwagDataset =
-    await loadHellaSwag(numDataPoints);
+async function evaluateModel(model: GPT | ONNXModel, numDataPoints = -1) {
+  const hellaswagDataset: HellaSwagDataset = await loadHellaSwag(numDataPoints);
   const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2");
   console.log("Starting the HellaSwag benchmark...");
 
   const start = Date.now();
-  const accuracy = await models.evaluate_hellaswag(
+  const accuracy = await evaluate_hellaswag(
     model,
     tokenizer,
     hellaswagDataset,
@@ -91,15 +94,15 @@ async function main(): Promise<void> {
     { helpArg: "help" },
   );
 
-  let model: models.GPT | models.ONNXModel | undefined;
+  let model: GPT | ONNXModel | undefined;
   switch (args.model) {
     case "onnx":
       console.log("Using ONNX pretrained model Xenova/gpt2");
-      model = await models.ONNXModel.init_pretrained("Xenova/gpt2");
+      model = await ONNXModel.init_pretrained("Xenova/gpt2");
       break;
     case "gpt-tfjs-random":
       console.log("Using GPT-TFJS with random initialization");
-      model = new models.GPT({ seed: 42 });
+      model = new GPT({ seed: 42 });
       break;
     case "gpt-tfjs-pretrained":
       console.log("Using GPT-TFJS with pretrained weights");
@@ -109,7 +112,7 @@ async function main(): Promise<void> {
         );
       }
       const encodedModel = await fsPromise.readFile(args.pretrainedModelPath);
-      model = (await serialization.model.decode(encodedModel)) as models.GPT;
+      model = (await modelDecode(encodedModel)) as GPT;
       break;
   }
   await evaluateModel(model, args.numDataPoints);

cli/src/train_gpt.ts

@@ -1,12 +1,13 @@
 import "@tensorflow/tfjs-node";
-import { models, Dataset, Tokenizer } from "@epfml/discojs";
+import type { GPTConfig } from "@epfml/discojs";
+import { GPT, Dataset, Tokenizer } from "@epfml/discojs";
 import { List } from "immutable";
 
 async function main(): Promise<void> {
   const data = "Lorem ipsum dolor sit amet, consectetur adipis";
   const seed = 42;
 
-  const config: models.GPTConfig = {
+  const config: GPTConfig = {
     modelType: "gpt-nano",
     lr: 0.01,
     maxIter: 50,
@@ -26,7 +27,7 @@ async function main(): Promise<void> {
     .repeat()
     .batch(8);
 
-  const model = new models.GPT(config);
+  const model = new GPT(config);
   for await (const logs of model.train(tokenDataset, undefined)) {
     console.log(logs);
   }

cli/src/user_log.ts

@@ -1,4 +1,5 @@
-import { args, BenchmarkArguments } from "./args.js";
+import type { BenchmarkArguments } from "./args.js";
+import { args } from "./args.js";
 import type { SummaryLogs, DataType, Network, Task } from "@epfml/discojs";
 
 type SerializableArguments = Omit<BenchmarkArguments, "provider" | "host"> & {

cli/tsconfig.json

@@ -1,14 +1,14 @@
 {
-  "extends": "../tsconfig.base.json",
+  "extends": "../tsconfig.base.lib.json",
   "references": [
     {
-      "path": "../discojs"
+      "path": "../discojs/tsconfig.lib.json"
     },
     {
-      "path": "../discojs-node"
+      "path": "../discojs-node/tsconfig.lib.json"
     },
     {
-      "path": "../server"
+      "path": "../server/tsconfig.lib.json"
     }
   ],
   "compilerOptions": {

discojs-node/package.json

@@ -2,7 +2,11 @@
   "name": "@epfml/discojs-node",
   "version": "3.0.0",
   "type": "module",
-  "exports": "./dist/index.js",
+  "exports": {
+    "@disco/source": "./src/index.ts",
+    "default": "./dist/index.js",
+    "types": "./dist/index.d.ts"
+  },
   "types": "dist/index.d.ts",
   "scripts": {
     "watch": "nodemon --ext ts --ignore dist --watch ../discojs/dist --watch . --exec pnpm run",