Merge pull request #314 from OpenMined/aziz/load_testing

load testing with locust

load_testing/locustfile.py

@@ -0,0 +1,33 @@
+from locust import FastHttpUser, between, task
+from loguru import logger
+
+from syftbox.client.plugins.sync import endpoints
+
+
+class SyftBoxUser(FastHttpUser):
+    network_timeout = 5.0
+    connection_timeout = 5.0
+    wait_time = between(0.5, 1.5)
+
+    def on_start(self):
+        self.datasites = []
+        self.email = "aziz@openmined.org"
+        self.remote_state: dict[str, list[endpoints.FileMetadata]] = {}
+
+    @task
+    def sync_datasites(self):
+        remote_datasite_states = endpoints.get_datasite_states(
+            self.client,
+            email=self.email,
+        )
+        # logger.info(f"Syncing {len(remote_datasite_states)} datasites")
+        all_files = []
+        for email, remote_state in remote_datasite_states.items():
+            all_files.extend(remote_state)
+
+        all_paths = [str(f.path) for f in all_files][:10]
+        logger.info(f"Downloading {len(all_paths)} files")
+        endpoints.download_bulk(
+            self.client,
+            all_paths,
+        )

pyproject.toml

@@ -50,6 +50,7 @@ dev-dependencies = [
     "bump2version>=1.0.1",
     "faker>=30.3.0",
     "ipykernel>=6.29.5",
+    "locust>=2.32.0",
     "pre-commit>=4.0.1",
     "pytest-cov>=5.0.0",
     "pytest-xdist[psutil]>=3.6.1",

syftbox/lib/prof.py

@@ -0,0 +1,47 @@
+# stdlib
+import contextlib
+import os
+import signal
+import subprocess  # nosec
+import tempfile
+import time
+
+
+@contextlib.contextmanager
+def pyspy() -> None:  # type: ignore
+    """Profile a block of code using py-spy. Intended for development purposes only.
+
+    Example:
+    ```
+    with pyspy():
+        # do some work
+        a = [i for i in range(1000000)]
+    ```
+    """
+    fd, fname = tempfile.mkstemp(".svg")
+    os.close(fd)
+
+    command = [
+        "sudo",
+        "-S",
+        "py-spy",
+        "record",
+        "-r",
+        "100",
+        "-o",
+        fname,
+        "--pid",
+        str(os.getpid()),
+    ]
+    process = subprocess.Popen(command, preexec_fn=os.setsid)  # nosec
+
+    start_time = time.time()
+    yield process
+    end_time = time.time()
+
+    print(f"Execution time: {end_time - start_time}")
+    try:
+        os.killpg(os.getpgid(process.pid), signal.SIGINT)
+        os.chmod(fname, 0o444)
+    except Exception as e:
+        print(f"Error: {e}")