WIP: Substra orchestrator

app.py

@@ -5,8 +5,12 @@ from substra_launcher import launch_substra_space
 api = HfApi()
 
 gr.Interface(
+    title="Susbtra Demo - Orchestrator",
     fn=lambda *args, **kwargs: launch_substra_space(api, *args, **kwargs),
-    inputs="text",
+    inputs=[
+        gr.Slider(label="num_hospitals", value=3, minimum=2, maximum=5, step=1),
+        gr.Text(label="space_name")
+    ],
     outputs="text",
-    examples=[["NimaBoscarino/substra-test"]]
+    examples=[[2, "substra-test"]]
 ).launch()

substra_launcher.py

@@ -1,11 +1,13 @@
 from huggingface_hub import HfApi, RepoUrl
 
 
-def launch_substra_space(hf_api: HfApi, repo_id: str) -> RepoUrl:
+def launch_substra_space(hf_api: HfApi, num_hospitals: int, repo_id: str) -> RepoUrl:
+    repo_id = "NimaBoscarino/" + repo_id
+
     repo_url = hf_api.create_repo(
         repo_id=repo_id,
         repo_type="space",
-        space_sdk="docker"
+        space_sdk="docker",
     )
 
     hf_api.upload_folder(
@@ -14,4 +16,15 @@ def launch_substra_space(hf_api: HfApi, repo_id: str) -> RepoUrl:
         folder_path="substra_template/"
     )
 
+    ENV_FILE = f"""\
+SUBSTRA_NUM_HOSPITALS={num_hospitals}
+    """
+
+    hf_api.upload_file(
+        repo_id=repo_id,
+        path_or_fileobj=ENV_FILE.encode(),
+        path_in_repo=".env",
+        repo_type="space",
+    )
+
     return repo_url

substra_template/Dockerfile

@@ -1,13 +1,15 @@
-FROM python:3.10-slim-bullseye
+FROM python:3.10
 
 # Set the working directory to /code
 WORKDIR /code
 
 # Copy the current directory contents into the container at /code
 COPY ./requirements.txt /code/requirements.txt
+COPY ./mlflow-2.1.2.dev0-py3-none-any.whl /code/mlflow-2.1.2.dev0-py3-none-any.whl
 
 # Install requirements.txt
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+RUN chmod -R 777 /usr/local/lib/python3.10/site-packages/
 
 # Set up a new user named "user" with user ID 1000
 RUN useradd -m -u 1000 user
@@ -23,5 +25,7 @@ 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
 
+RUN chmod -R 777 $HOME/app/
+
 EXPOSE 7860
-CMD ["bash", "-c", "/code/run.sh"]
+CMD ["bash", "run.sh"]

substra_template/app.py

@@ -1,13 +0,0 @@
-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

@@ -6,7 +6,7 @@ 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
+TIMEOUT = 240  # 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

substra_template/requirements.txt

@@ -9,4 +9,5 @@ Pillow
 transformers
 matplotlib
 pandas
-mlflow
+python-dotenv
+./mlflow-2.1.2.dev0-py3-none-any.whl

substra_template/run.sh

@@ -1,3 +1,13 @@
-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
+PYTHONPATH=$HOME/app python run_compute_plan.py &
+PYTHONPATH=$HOME/app python mlflow_live_performances.py &
+
+SITE_PACKAGES=$(python -c 'import site; print(site.getsitepackages()[0])')
+
+# Fix for the UI code being embedded in an iframe
+# Replace window.parent.location.origin with *
+for i in $SITE_PACKAGES/mlflow/server/js/build/static/js/*.js; do
+  sed -i 's/window\.parent\.location\.origin)/"*")/' $i
+  sed 's/window.top?.location.href || window.location.href/window.location.href/g' -i $i
+done
+
+mlflow ui --port 7860 --host 0.0.0.0

substra_template/run_compute_plan.py

@@ -5,13 +5,19 @@ from substrafl.strategies import FedAvg
 
 import torch
 
+from dotenv import load_dotenv
+import os
+load_dotenv()
+
+NUM_CLIENTS = int(os.environ["NUM_CLIENTS"])
+
 seed = 42
 torch.manual_seed(seed)
 model = CNN()
 optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 criterion = torch.nn.CrossEntropyLoss()
 
-runner = SubstraRunner()
+runner = SubstraRunner(num_clients=NUM_CLIENTS)
 runner.set_up_clients()
 runner.prepare_data()
 runner.register_data()

substra_template/substra_helpers/substra_runner.py

@@ -30,8 +30,8 @@ import torch
 
 
 class SubstraRunner:
-    def __init__(self):
-        self.num_clients = 3
+    def __init__(self, num_clients: int):
+        self.num_clients = num_clients
         self.clients = {}
         self.algo_provider: Optional[Client] = None