args.ts

import { parse } from "ts-command-line-args";
import { Map, Set } from "immutable";

import type { DataType, Network, TaskProvider } from "@epfml/discojs";
import { defaultTasks } from "@epfml/discojs";

type AggregationStrategy = "mean" | "byzantine" | "secure";

function parseAggregator(raw: string): AggregationStrategy {
  if (raw === "mean" || raw == "byzantine" || raw == "secure") return raw;
  else throw new Error(`Aggregator ${raw} is not supported.`);
}

export interface BenchmarkArguments {
  provider: TaskProvider<DataType, Network>;
  testID: string;
  numberOfUsers: number;
  epochs: number;
  roundDuration: number;
  batchSize: number;
  validationSplit: number;

  // DP
  epsilon?: number;
  delta?: number;
  dpDefaultClippingRadius?: number;
  // Aggregator
  aggregator: AggregationStrategy;
  // Byzantine aggregator
  clippingRadius?: number;
  maxIterations?: number;
  beta?: number;
  // Secure aggregator
  maxShareValue?: number;

  save: boolean;
  host: URL;
}

type BenchmarkUnsafeArguments = Omit<BenchmarkArguments, "provider"> & {
  task: string;
  help?: boolean;
};

const argExample = "e.g. npm start -- -u 2 -e 3 # runs 2 users for 3 epochs";

const unsafeArgs = parse<BenchmarkUnsafeArguments>(
  {
    testID: { type: String, alias: "i", description: "ID of the testcase" },
    task: {
      type: String,
      alias: "t",
      description:
        "Task: tinder_dog, titanic, simple_face, cifar10 or lus_covid",
      defaultValue: "tinder_dog",
    },
    numberOfUsers: {
      type: Number,
      alias: "u",
      description: "Number of users",
      defaultValue: 2,
    },
    epochs: {
      type: Number,
      alias: "e",
      description: "Number of epochs",
      defaultValue: 10,
    },
    roundDuration: {
      type: Number,
      alias: "r",
      description: "Round duration (in epochs)",
      defaultValue: 2,
    },
    batchSize: {
      type: Number,
      alias: "b",
      description: "Training batch size",
      defaultValue: 10,
    },
    validationSplit: {
      type: Number,
      alias: "v",
      description: "Validation dataset ratio",
      defaultValue: 0.2,
    },
    save: {
      type: Boolean,
      alias: "s",
      description: "Save logs of benchmark",
      defaultValue: false,
    },
    host: {
      type: (raw: string) => new URL(raw),
      typeLabel: "URL",
      description: "Host to connect to",
      defaultValue: new URL("http://localhost:8080"),
    },

    // Aggregator
    aggregator: {
      type: parseAggregator,
      description: "Type of weight aggregator",
      defaultValue: "mean",
    },

    // Byzantine aggregator
    clippingRadius: {
      type: Number,
      description: "Clipping radius for centered clipping",
      optional: true,
    },
    maxIterations: {
      type: Number,
      description: "Maximum centered clipping iterations",
      optional: true,
    },
    beta: {
      type: Number,
      description:
        "Momentum coefficient to smooth the aggregation over multiple rounds",
      optional: true,
    },

    // Secure aggregator
    maxShareValue: {
      type: Number,
      description: "Maximum absolute value over all the weights",
      optional: true,
    },

    // Differential Privacy
    epsilon: {
      type: Number,
      description: "Privacy budget",
      optional: true,
      defaultValue: undefined,
    },
    delta: {
      type: Number,
      description: "Probability of failure, slack parameter",
      optional: true,
      defaultValue: undefined,
    },
    dpDefaultClippingRadius: {
      type: Number,
      description: "Default clipping radius for DP",
      optional: true,
      defaultValue: undefined,
    },

    help: {
      type: Boolean,
      optional: true,
      alias: "h",
      description: "Prints this usage guide",
    },
  },
  {
    helpArg: "help",
    headerContentSections: [
      { header: "DISCO CLI", content: "npm start -- [Options]\n" + argExample },
    ],
  },
);

const supportedTasks = Map(
  await Promise.all(
    Set.of<TaskProvider<"image" | "tabular", Network>>(
      defaultTasks.cifar10,
      defaultTasks.lusCovid,
      defaultTasks.simpleFace,
      defaultTasks.titanic,
      defaultTasks.tinderDog,
      defaultTasks.mnist,
    ).map(
      async (t) =>
        [(await t.getTask()).id, t] as [
          string,
          TaskProvider<"image" | "tabular", Network>,
        ],
    ),
  ),
);

const provider = supportedTasks.get(unsafeArgs.task);
if (provider === undefined) {
  throw Error(`${unsafeArgs.task} not implemented.`);
}

export const args: BenchmarkArguments = {
  ...unsafeArgs,
  provider: {
    async getTask() {
      const task = await provider.getTask();

      // Override training information
      task.trainingInformation.batchSize = unsafeArgs.batchSize;
      task.trainingInformation.roundDuration = unsafeArgs.roundDuration;
      task.trainingInformation.epochs = unsafeArgs.epochs;
      task.trainingInformation.validationSplit = unsafeArgs.validationSplit;

      const { aggregator, clippingRadius, maxIterations, beta, maxShareValue } =
        unsafeArgs;

      // For aggregators
      if (aggregator !== undefined)
        task.trainingInformation.aggregationStrategy = aggregator;

      // For byzantine aggregator
      if (
        clippingRadius !== undefined &&
        maxIterations !== undefined &&
        beta !== undefined
      ) {
        if (task.trainingInformation.scheme === "local")
          throw new Error(
            "Byzantine aggregator is not supported for local training",
          );
        if (task.trainingInformation.aggregationStrategy !== "byzantine")
          throw new Error(
            "Byzantine parameters can be set only when aggregationStrategy is byzantine",
          );

        task.trainingInformation.privacy = {
          ...task.trainingInformation.privacy,
          byzantineFaultTolerance: {
            clippingRadius,
            maxIterations,
            beta,
          },
        };
      }

      // For secure aggregator
      if (maxShareValue !== undefined) {
        if (task.trainingInformation.scheme !== "decentralized")
          throw new Error(
            "Secure aggation is only supported for decentralized laerning",
          );
        if (task.trainingInformation.aggregationStrategy !== "secure")
          throw new Error(
            "maxShareValue can be set when aggregationStrategy is secure",
          );

        task.trainingInformation.maxShareValue = maxShareValue;
      }

      // For DP
      const { dpDefaultClippingRadius, epsilon, delta } = unsafeArgs;

      if (
        // dpDefaultClippingRadius !== undefined &&
        epsilon !== undefined &&
        delta !== undefined
      ) {
        if (task.trainingInformation.scheme === "local")
          throw new Error("Can't have differential privacy for local training");

        const defaultRadius = dpDefaultClippingRadius
          ? dpDefaultClippingRadius
          : 1;

        // for the case where privacy parameters are not defined in the default tasks
        task.trainingInformation.privacy ??= {};
        task.trainingInformation.privacy.differentialPrivacy = {
          clippingRadius: defaultRadius,
          epsilon: epsilon,
          delta: delta,
        };
      }

      return task;
    },
    getModel: () => provider.getModel(),
  },
};