WIP: Substra orchestrator

app.py

@@ -0,0 +1,12 @@
+import gradio as gr
+from huggingface_hub import HfApi
+from substra_launcher import launch_substra_space
+
+api = HfApi()
+
+gr.Interface(
+    fn=lambda *args, **kwargs: launch_substra_space(api, *args, **kwargs),
+    inputs="text",
+    outputs="text",
+    examples=[["NimaBoscarino/substra-test"]]
+).launch()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+pytest
+huggingface_hub

substra_launcher.py

@@ -0,0 +1,17 @@
+from huggingface_hub import HfApi, RepoUrl
+
+
+def launch_substra_space(hf_api: HfApi, repo_id: str) -> RepoUrl:
+    repo_url = hf_api.create_repo(
+        repo_id=repo_id,
+        repo_type="space",
+        space_sdk="docker"
+    )
+
+    hf_api.upload_folder(
+        repo_id=repo_id,
+        repo_type="space",
+        folder_path="substra_template/"
+    )
+
+    return repo_url

substra_template/Dockerfile

@@ -0,0 +1,27 @@
+FROM python:3.10-slim-bullseye
+
+# Set the working directory to /code
+WORKDIR /code
+
+# Copy the current directory contents into the container at /code
+COPY ./requirements.txt /code/requirements.txt
+
+# Install requirements.txt
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+# Switch to the "user" user
+USER user
+# Set home to the user's home directory
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+EXPOSE 7860
+CMD ["bash", "-c", "/code/run.sh"]

substra_template/app.py

@@ -0,0 +1,13 @@
+import gradio as gr
+
+
+def read_logs():
+    with open("output.log", "r") as f:
+        return f.read()
+
+
+with gr.Blocks() as demo:
+    logs = gr.Plot()
+    demo.load(read_logs, None, logs, every=1)
+
+demo.queue().launch()

substra_template/mlflow_live_performances.py

@@ -0,0 +1,45 @@
+import pandas as pd
+import json
+from pathlib import Path
+from mlflow import log_metric
+import time
+import os
+from glob import glob
+
+TIMEOUT = 60  # Number of seconds to stop the script after the last update of the json file
+POLLING_FREQUENCY = 10  # Try to read the updates in the file every 10 seconds
+
+# Wait for the file to be found
+start = time.time()
+while not len(glob(str(Path("local-worker") / "live_performances" / "*" / "performances.json"))) > 0:
+    time.sleep(POLLING_FREQUENCY)
+    if time.time() - start >= TIMEOUT:
+        raise TimeoutError("The performance file does not exist, maybe no test task has been executed yet.")
+
+path_to_json = Path(glob(str(Path("local-worker") / "live_performances" / "*" / "performances.json"))[0])
+
+logged_rows = []
+last_update = time.time()
+
+while (time.time() - last_update) <= TIMEOUT:
+
+    if last_update == os.path.getmtime(str(path_to_json)):
+        time.sleep(POLLING_FREQUENCY)
+        continue
+
+    last_update = os.path.getmtime(str(path_to_json))
+
+    time.sleep(1)  # Waiting for the json to be fully written
+    dict_perf = json.load(path_to_json.open())
+
+    df = pd.DataFrame(dict_perf)
+
+    for _, row in df.iterrows():
+        if row["testtask_key"] in logged_rows:
+            continue
+
+        logged_rows.append(row["testtask_key"])
+
+        step = int(row["round_idx"]) if row["round_idx"] is not None else int(row["testtask_rank"])
+
+        log_metric(f"{row['metric_name']}_{row['worker']}", row["performance"], step)

substra_template/requirements.txt

@@ -0,0 +1,12 @@
+gradio
+substrafl
+datasets
+torch
+torchvision
+scikit-learn
+numpy==1.23.0
+Pillow
+transformers
+matplotlib
+pandas
+mlflow

substra_template/run.sh

@@ -0,0 +1,3 @@
+PYTHONPATH=/Users/nima/Work/society-ethics/substra/substra_template python run_compute_plan.py &
+PYTHONPATH=/Users/nima/Work/society-ethics/substra/substra_template python mlflow_live_performances.py &
+mlflow ui

substra_template/run_compute_plan.py

@@ -0,0 +1,34 @@
+from substra_helpers.substra_runner import SubstraRunner, algo_generator
+from substra_helpers.model import CNN
+from substra_helpers.dataset import TorchDataset
+from substrafl.strategies import FedAvg
+
+import torch
+
+seed = 42
+torch.manual_seed(seed)
+model = CNN()
+optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+criterion = torch.nn.CrossEntropyLoss()
+
+runner = SubstraRunner()
+runner.set_up_clients()
+runner.prepare_data()
+runner.register_data()
+runner.register_metric()
+
+runner.algorithm = algo_generator(
+    model=model,
+    criterion=criterion,
+    optimizer=optimizer,
+    index_generator=runner.index_generator,
+    dataset=TorchDataset,
+    seed=seed
+)()
+
+runner.strategy = FedAvg()
+
+runner.set_aggregation()
+runner.set_testing()
+
+runner.run_compute_plan()

substra_template/substra_helpers/dataset.py

@@ -0,0 +1,29 @@
+import torch
+from torch.utils import data
+import torch.nn.functional as F
+import numpy as np
+
+
+class TorchDataset(data.Dataset):
+    def __init__(self, datasamples, is_inference: bool):
+        self.x = datasamples["image"]
+        self.y = datasamples["label"]
+        self.is_inference = is_inference
+
+    def __getitem__(self, idx):
+
+        if self.is_inference:
+            x = torch.FloatTensor(np.array(self.x[idx])[None, ...]) / 255
+            return x
+
+        else:
+            x = torch.FloatTensor(np.array(self.x[idx])[None, ...]) / 255
+
+            y = torch.tensor(self.y[idx]).type(torch.int64)
+            y = F.one_hot(y, 10)
+            y = y.type(torch.float32)
+
+            return x, y
+
+    def __len__(self):
+        return len(self.x)

substra_template/substra_helpers/dataset_assets/description.md

@@ -0,0 +1,18 @@
+# Mnist
+
+This dataset is [THE MNIST DATABASE of handwritten digits](http://yann.lecun.com/exdb/mnist/).
+
+The target is the number (0 -> 9) represented by the pixels.
+
+## Data repartition
+
+### Train and test
+
+### Split data between organizations
+
+## Opener usage
+
+The opener exposes 2 methods:
+
+- `get_data` returns a dictionary containing the images and the labels as numpy arrays
+- `fake_data` returns a fake data sample of images and labels in a dict

substra_template/substra_helpers/dataset_assets/opener.py

@@ -0,0 +1,20 @@
+import numpy as np
+import substratools as tools
+from datasets import load_from_disk
+from transformers import ImageFeatureExtractionMixin
+
+
+class MnistOpener(tools.Opener):
+    def fake_data(self, n_samples=None):
+        N_SAMPLES = n_samples if n_samples and n_samples <= 100 else 100
+
+        fake_images = np.random.randint(256, size=(N_SAMPLES, 28, 28))
+
+        fake_labels = np.random.randint(10, size=N_SAMPLES)
+
+        data = {"image": fake_images, "label": fake_labels}
+
+        return data
+
+    def get_data(self, folders):
+        return load_from_disk(folders[0])

substra_template/substra_helpers/model.py

@@ -0,0 +1,25 @@
+from torch import nn
+import torch.nn.functional as F
+
+
+# TODO: Would be cool to use a simple Transformer model... then I could use the Trainer API 👀
+class CNN(nn.Module):
+    def __init__(self):
+        super(CNN, self).__init__()
+        self.conv1 = nn.Conv2d(1, 32, kernel_size=5)
+        self.conv2 = nn.Conv2d(32, 32, kernel_size=5)
+        self.conv3 = nn.Conv2d(32, 64, kernel_size=5)
+        self.fc1 = nn.Linear(3 * 3 * 64, 256)
+        self.fc2 = nn.Linear(256, 10)
+
+    def forward(self, x, eval=False):
+        x = F.relu(self.conv1(x))
+        x = F.relu(F.max_pool2d(self.conv2(x), 2))
+        x = F.dropout(x, p=0.5, training=not eval)
+        x = F.relu(F.max_pool2d(self.conv3(x), 2))
+        x = F.dropout(x, p=0.5, training=not eval)
+        x = x.view(-1, 3 * 3 * 64)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, p=0.5, training=not eval)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)

substra_template/substra_helpers/substra_runner.py

@@ -0,0 +1,194 @@
+import pathlib
+import shutil
+from typing import Optional, List
+
+from substra import Client, BackendType
+
+from substra.sdk.schemas import (
+    DatasetSpec,
+    Permissions,
+    DataSampleSpec
+)
+
+from substrafl.strategies import Strategy
+from substrafl.dependency import Dependency
+from substrafl.remote.register import add_metric
+from substrafl.index_generator import NpIndexGenerator
+from substrafl.algorithms.pytorch import TorchFedAvgAlgo
+
+from substrafl.nodes import TrainDataNode, AggregationNode, TestDataNode
+from substrafl.evaluation_strategy import EvaluationStrategy
+
+from substrafl.experiment import execute_experiment
+from substra.sdk.models import ComputePlan
+
+from datasets import load_dataset, Dataset
+from sklearn.metrics import accuracy_score
+import numpy as np
+
+import torch
+
+
+class SubstraRunner:
+    def __init__(self):
+        self.num_clients = 3
+        self.clients = {}
+        self.algo_provider: Optional[Client] = None
+
+        self.datasets: List[Dataset] = []
+        self.test_dataset: Optional[Dataset] = None
+        self.path = pathlib.Path(__file__).parent.resolve()
+
+        self.dataset_keys = {}
+        self.train_data_sample_keys = {}
+        self.test_data_sample_keys = {}
+
+        self.metric_key: Optional[str] = None
+
+        NUM_UPDATES = 100
+        BATCH_SIZE = 32
+
+        self.index_generator = NpIndexGenerator(
+            batch_size=BATCH_SIZE,
+            num_updates=NUM_UPDATES,
+        )
+
+        self.algorithm: Optional[TorchFedAvgAlgo] = None
+        self.strategy: Optional[Strategy] = None
+
+        self.aggregation_node: Optional[AggregationNode] = None
+        self.train_data_nodes = list()
+        self.test_data_nodes = list()
+        self.eval_strategy: Optional[EvaluationStrategy] = None
+
+        self.NUM_ROUNDS = 3
+        self.compute_plan: Optional[ComputePlan] = None
+
+        self.experiment_folder = self.path / "experiment_summaries"
+
+    def set_up_clients(self):
+        self.algo_provider = Client(backend_type=BackendType.LOCAL_SUBPROCESS)
+
+        self.clients = {
+            c.organization_info().organization_id: c
+            for c in [Client(backend_type=BackendType.LOCAL_SUBPROCESS) for _ in range(self.num_clients - 1)]
+        }
+
+    def prepare_data(self):
+        dataset = load_dataset("mnist", split="train").shuffle()
+        self.datasets = [dataset.shard(num_shards=self.num_clients - 1, index=i) for i in range(self.num_clients - 1)]
+
+        self.test_dataset = load_dataset("mnist", split="test")
+
+        data_path = self.path / "data"
+        if data_path.exists() and data_path.is_dir():
+            shutil.rmtree(data_path)
+
+        for i, client_id in enumerate(self.clients):
+            ds = self.datasets[i]
+            ds.save_to_disk(data_path / client_id / "train")
+            self.test_dataset.save_to_disk(data_path / client_id / "test")
+
+    def register_data(self):
+        for client_id, client in self.clients.items():
+            permissions_dataset = Permissions(public=False, authorized_ids=[
+                self.algo_provider.organization_info().organization_id
+            ])
+
+            dataset = DatasetSpec(
+                name="MNIST",
+                type="npy",
+                data_opener=self.path / pathlib.Path("dataset_assets/opener.py"),
+                description=self.path / pathlib.Path("dataset_assets/description.md"),
+                permissions=permissions_dataset,
+                logs_permission=permissions_dataset,
+            )
+            self.dataset_keys[client_id] = client.add_dataset(dataset)
+            assert self.dataset_keys[client_id], "Missing dataset key"
+
+            self.train_data_sample_keys[client_id] = client.add_data_sample(DataSampleSpec(
+                data_manager_keys=[self.dataset_keys[client_id]],
+                path=self.path / "data" / client_id / "train",
+            ))
+
+            data_sample = DataSampleSpec(
+                data_manager_keys=[self.dataset_keys[client_id]],
+                path=self.path / "data" / client_id / "test",
+            )
+            self.test_data_sample_keys[client_id] = client.add_data_sample(data_sample)
+
+    def register_metric(self):
+        permissions_metric = Permissions(
+                                public=False,
+                                authorized_ids=[
+                                   self.algo_provider.organization_info().organization_id
+                                ] + list(self.clients.keys())
+                            )
+
+        metric_deps = Dependency(pypi_dependencies=["numpy==1.23.1", "scikit-learn==1.1.1"])
+
+        def accuracy(datasamples, predictions_path):
+            y_true = datasamples["label"]
+            y_pred = np.load(predictions_path)
+
+            return accuracy_score(y_true, np.argmax(y_pred, axis=1))
+
+        self.metric_key = add_metric(
+            client=self.algo_provider,
+            metric_function=accuracy,
+            permissions=permissions_metric,
+            dependencies=metric_deps,
+        )
+
+    def set_aggregation(self):
+        self.aggregation_node = AggregationNode(self.algo_provider.organization_info().organization_id)
+
+        for org_id in self.clients:
+            train_data_node = TrainDataNode(
+                organization_id=org_id,
+                data_manager_key=self.dataset_keys[org_id],
+                data_sample_keys=[self.train_data_sample_keys[org_id]],
+            )
+            self.train_data_nodes.append(train_data_node)
+
+    def set_testing(self):
+        for org_id in self.clients:
+            test_data_node = TestDataNode(
+                organization_id=org_id,
+                data_manager_key=self.dataset_keys[org_id],
+                test_data_sample_keys=[self.test_data_sample_keys[org_id]],
+                metric_keys=[self.metric_key],
+            )
+            self.test_data_nodes.append(test_data_node)
+
+        self.eval_strategy = EvaluationStrategy(test_data_nodes=self.test_data_nodes, rounds=1)
+
+    def run_compute_plan(self):
+        algo_deps = Dependency(pypi_dependencies=["numpy==1.23.1", "torch==1.11.0"])
+
+        self.compute_plan = execute_experiment(
+            client=self.algo_provider,
+            algo=self.algorithm,
+            strategy=self.strategy,
+            train_data_nodes=self.train_data_nodes,
+            evaluation_strategy=self.eval_strategy,
+            aggregation_node=self.aggregation_node,
+            num_rounds=self.NUM_ROUNDS,
+            experiment_folder=self.experiment_folder,
+            dependencies=algo_deps,
+        )
+
+
+def algo_generator(model, criterion, optimizer, index_generator, dataset, seed):
+    class MyAlgo(TorchFedAvgAlgo):
+        def __init__(self):
+            super().__init__(
+                model=model,
+                criterion=criterion,
+                optimizer=optimizer,
+                index_generator=index_generator,
+                dataset=dataset,
+                seed=seed,
+            )
+
+    return MyAlgo

tests/test_substra_launcher.py

@@ -0,0 +1,25 @@
+import pytest
+from unittest.mock import Mock
+
+from substra_launcher import launch_substra_space
+
+
+class TestSubstraLauncher:
+
+    @pytest.fixture
+    def mock_hf_api(self):
+        mock_hf_api = Mock()
+        mock_hf_api.create_repo = Mock(side_effect=lambda repo_id, *args, **kwargs: f"https://hf.space/{repo_id}")
+        return mock_hf_api
+
+    def test_launch_substra_space(self, mock_hf_api):
+        repo_id = "user/space"
+        repo_link = launch_substra_space(mock_hf_api, repo_id=repo_id)
+        mock_hf_api.create_repo.assert_called_once_with(
+            repo_id=repo_id, repo_type="space", space_sdk="docker"
+        )
+        mock_hf_api.upload_folder.assert_called_once_with(
+            repo_id=repo_id, repo_type="space", folder_path="substra_template/"
+        )
+
+        assert repo_link == f"https://hf.space/{repo_id}"

tests/test_substra_runner.py

@@ -0,0 +1,55 @@
+import pytest
+from unittest.mock import Mock, call
+from datasets import Dataset
+
+from substra_template.substra_runner import SubstraRunner
+
+
+class TestSubstraRunner:
+    @pytest.fixture
+    def mock_substra_client_class(self, monkeypatch):
+        mock_substra_client_class = Mock()
+        monkeypatch.setattr("substra_template.substra_runner.Client", mock_substra_client_class)
+
+        return mock_substra_client_class
+
+    @pytest.fixture
+    def mock_load_dataset(self, monkeypatch):
+        mock_load_dataset = Mock()
+        monkeypatch.setattr("substra_template.substra_runner.load_dataset", mock_load_dataset)
+
+        return mock_load_dataset
+
+    def test_set_up_clients(self, mock_substra_client_class):
+        runner = SubstraRunner()
+        runner.set_up_clients()
+
+        mock_substra_client_class.assert_called()
+
+    def test_prepare_data(self, mock_load_dataset):
+        runner = SubstraRunner()
+        runner.prepare_data()
+
+        mock_load_dataset.assert_has_calls(calls=[
+            call("mnist", split="train"),
+            call("mnist", split="test"),
+        ], any_order=True)
+
+        assert len(runner.datasets) == runner.num_clients - 1
+
+    def test_register_data(self, mock_load_dataset):
+        runner = SubstraRunner()
+        runner.datasets = [Dataset.from_dict({}) for _ in range(runner.num_clients - 1)]
+
+        runner.register_data()
+
+    def test_register_metric(self):
+        runner = SubstraRunner()
+        runner.set_up_clients()
+        runner.register_metric()
+
+    def test_set_aggregation(self):
+        pass
+
+    def test_set_testing(self):
+        pass