diff --git a/notebooks/enclave/double-blind-eval-w-receipts/1. DO-benchmark-owner-dbe.ipynb b/notebooks/enclave/double-blind-eval-w-receipts/1. DO-benchmark-owner-dbe.ipynb index 2d6973ebcfa..a62cc6659ec 100644 --- a/notebooks/enclave/double-blind-eval-w-receipts/1. DO-benchmark-owner-dbe.ipynb +++ b/notebooks/enclave/double-blind-eval-w-receipts/1. DO-benchmark-owner-dbe.ipynb @@ -150,8 +150,8 @@ "source": [ "### Step 1.1 — Wait until both accept\n", "\n", - "The enclave accepts automatically; the model owner approves us from their notebook. Re-run the cell\n", - "below until both show as accepted." + "The enclave accepts automatically; the model owner approves us from their notebook. The cell below\n", + "shows both peers, and the cell after it waits until the model owner has approved." ] }, { @@ -169,11 +169,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "benchmark_owner.validate_peer(MODEL_OWNER_EMAIL)\n", - "print(f\" Connected to {MODEL_OWNER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [MODEL_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " benchmark_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -216,14 +216,11 @@ " expected_syft_version=None,\n", ")\n", "\n", - "result = benchmark_owner.attest_peer(ENCLAVE_EMAIL, policy=policy)\n", - "\n", - "if result is None:\n", - " print(\" 🟠 The enclave published no attestation yet. Wait a moment and re-run this cell.\")\n", - "else:\n", - " result.print_checklist()\n", - " print()\n", - " print(f\" Bound to the report: {result.verified_key_bundle is not None}\")" + "# Wait until the enclave publishes its attestation, then verify it.\n", + "result = benchmark_owner.wait_until_attested(ENCLAVE_EMAIL, policy=policy)\n", + "result.print_checklist()\n", + "print()\n", + "print(f\" Bound to the report: {result.verified_key_bundle is not None}\")" ] }, { @@ -430,11 +427,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Unlike `jobs`, `datasets` does not sync when you read it, so pull first —\n", - "# the model owner may have uploaded since our last sync.\n", - "benchmark_owner.sync()\n", - "\n", - "model_dataset = benchmark_owner.datasets.get(\"dbe_adapter\", datasite=MODEL_OWNER_EMAIL)\n", + "# Wait until the model owner's card is here.\n", + "model_dataset = benchmark_owner.wait_until_has_dataset(\"dbe_adapter\", datasite=MODEL_OWNER_EMAIL)\n", "print(model_dataset.mock_files[0].read_text())" ] }, @@ -954,14 +948,11 @@ "metadata": {}, "outputs": [], "source": [ - "job = next((j for j in benchmark_owner.jobs if j.name == \"dbe_eval_job\"), None)\n", - "\n", - "if job is not None and not job.can_approve:\n", - " print(\" 🟠 Job 'dbe_eval_job' is here, but the enclave has not sent our approval file yet. Wait a moment and re-run this cell.\")\n", - "elif job is None:\n", - " print(\" 🟠 Job 'dbe_eval_job' not visible yet — the enclave is still distributing it. Wait a moment and re-run this cell.\")\n", - "else:\n", - " print(f\" ✅ Benchmark owner sees 'dbe_eval_job' status={job.status}\")" + "# Wait until the enclave has distributed the job and sent us our approval file.\n", + "job = benchmark_owner.wait_until_has_job(\n", + " \"dbe_eval_job\", where=lambda j: j.can_approve, timeout=900\n", + ")\n", + "print(f\" ✅ Benchmark owner sees 'dbe_eval_job' status={job.status}\")" ] }, { @@ -998,7 +989,8 @@ "\n", "Once the model owner approves, the enclave runs the evaluation and sends the output back to us. It\n", "installs PyTorch, downloads the base model and generates on two CPUs, which took about three\n", - "minutes on the release this was tested against. Re-run until the status is `done`." + "minutes on the release this was tested against. The cell below waits until the status is `done`.\n", + "If the job fails, the cell stops with an error: see 'If the job failed' below." ] }, { @@ -1007,15 +999,16 @@ "metadata": {}, "outputs": [], "source": [ - "job = next(j for j in benchmark_owner.jobs if j.name == \"dbe_eval_job\")\n", - "print(f\" Job status : {job.status}\")\n", + "from syft_rds.waiting import JobEndedError\n", "\n", - "if job.status == \"failed\":\n", - " print(\" ❌ The job failed inside the enclave. See 'If the job failed' below to read its logs.\")\n", - "elif job.status != \"done\" or not job.output_paths:\n", - " print(\" 🟠 Not finished yet — wait until the status is 'done', then re-run this cell.\")\n", - "else:\n", - " print(f\" ✅ Output files : {[p.name for p in job.output_paths]}\")" + "try:\n", + " job = benchmark_owner.wait_until_has_job(\n", + " \"dbe_eval_job\", status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + " )\n", + "except JobEndedError:\n", + " print(\" ❌ The job did not finish in the enclave. See 'If the job failed' below to read its logs.\")\n", + " raise\n", + "print(f\" ✅ Output files : {[p.name for p in job.output_paths]}\")" ] }, { diff --git a/notebooks/enclave/double-blind-eval-w-receipts/2. DO-model-owner-dbe.ipynb b/notebooks/enclave/double-blind-eval-w-receipts/2. DO-model-owner-dbe.ipynb index 2be6b118872..9dcc0244ad4 100644 --- a/notebooks/enclave/double-blind-eval-w-receipts/2. DO-model-owner-dbe.ipynb +++ b/notebooks/enclave/double-blind-eval-w-receipts/2. DO-model-owner-dbe.ipynb @@ -147,8 +147,8 @@ "source": [ "### Step 1.1 — Approve the benchmark owner\n", "\n", - "Re-run the cell below until their request appears, then approve it. `peer_must_exist=False` lets us\n", - "approve a request that has arrived before our own view of them has caught up." + "The cell below shows our peers. `peer_must_exist=False` lets us approve the benchmark owner before\n", + "their request reaches us, and the cell after it waits until the connection is live." ] }, { @@ -176,11 +176,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "model_owner.validate_peer(BENCHMARK_OWNER_EMAIL)\n", - "print(f\" Connected to {BENCHMARK_OWNER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [BENCHMARK_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " model_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -216,14 +216,11 @@ " expected_syft_version=None,\n", ")\n", "\n", - "result = model_owner.attest_peer(ENCLAVE_EMAIL, policy=policy)\n", - "\n", - "if result is None:\n", - " print(\" 🟠 The enclave published no attestation yet. Wait a moment and re-run this cell.\")\n", - "else:\n", - " result.print_checklist()\n", - " print()\n", - " print(f\" Bound to the report: {result.verified_key_bundle is not None}\")" + "# Wait until the enclave publishes its attestation, then verify it.\n", + "result = model_owner.wait_until_attested(ENCLAVE_EMAIL, policy=policy)\n", + "result.print_checklist()\n", + "print()\n", + "print(f\" Bound to the report: {result.verified_key_bundle is not None}\")" ] }, { @@ -362,7 +359,7 @@ "## Step 5 — Wait for the evaluation job\n", "\n", "The benchmark owner submits the job that runs our model on their prompts. The enclave distributes it\n", - "to both of us for approval. Re-run the cell below until it appears." + "to both of us for approval. The cell below waits until it is here and we can approve it." ] }, { @@ -371,14 +368,11 @@ "metadata": {}, "outputs": [], "source": [ - "job = next((j for j in model_owner.jobs if j.name == \"dbe_eval_job\"), None)\n", - "\n", - "if job is not None and not job.can_approve:\n", - " print(\" 🟠 Job 'dbe_eval_job' is here, but the enclave has not sent our approval file yet. Wait a moment and re-run this cell.\")\n", - "elif job is None:\n", - " print(\" 🟠 Job 'dbe_eval_job' not visible yet — the benchmark owner may not have submitted it. Wait a moment and re-run this cell.\")\n", - "else:\n", - " print(f\" ✅ Model owner sees 'dbe_eval_job' status={job.status}\")" + "# Wait until the enclave has distributed the job and sent us our approval file.\n", + "job = model_owner.wait_until_has_job(\n", + " \"dbe_eval_job\", where=lambda j: j.can_approve, timeout=900\n", + ")\n", + "print(f\" ✅ Model owner sees 'dbe_eval_job' status={job.status}\")" ] }, { diff --git a/notebooks/enclave/gemma/colab-3-persona/1. DO-model-owner-gemma-restrict.ipynb b/notebooks/enclave/gemma/colab-3-persona/1. DO-model-owner-gemma-restrict.ipynb index c0363f3b148..a2dcf90cc9d 100644 --- a/notebooks/enclave/gemma/colab-3-persona/1. DO-model-owner-gemma-restrict.ipynb +++ b/notebooks/enclave/gemma/colab-3-persona/1. DO-model-owner-gemma-restrict.ipynb @@ -165,7 +165,7 @@ "source": [ "### Step 1.1 — Wait for the Researcher peer request, then approve\n", "\n", - "The Researcher notebook adds you as a peer. Re-run the cell below until you see their request appear, then approve." + "The Researcher notebook adds you as a peer. The cell below shows your peers. `peer_must_exist=False` lets you approve before their request arrives, and the next cell waits until the connection is live." ] }, { @@ -198,11 +198,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "model_owner.validate_peer(RESEARCHER_EMAIL)\n", - "print(f\" Connected to {RESEARCHER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [RESEARCHER_EMAIL, ENCLAVE_EMAIL]:\n", + " model_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -222,8 +222,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", - "model_owner.attest_peer(ENCLAVE_EMAIL)" + "# Wait until the enclave publishes its attestation, then verify it.\n", + "model_owner.wait_until_attested(ENCLAVE_EMAIL)" ] }, { @@ -492,8 +492,7 @@ "metadata": {}, "outputs": [], "source": [ - "model_owner.sync()\n", - "restrict_job = next(j for j in model_owner.jobs if j.name == RESTRICT_JOB_NAME)\n", + "restrict_job = model_owner.wait_until_has_job(RESTRICT_JOB_NAME, where=lambda j: j.can_approve)\n", "print(f\" Model owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" ] }, @@ -537,8 +536,7 @@ "outputs": [], "source": [ "JOB_NAME = \"safety_eval_job\"\n", - "model_owner.sync()\n", - "model_owner_job = next(j for j in model_owner.jobs if j.name == JOB_NAME)\n", + "model_owner_job = model_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve, timeout=900)\n", "print(f\" Model owner sees '{JOB_NAME}' status={model_owner_job.status}\")" ] }, diff --git a/notebooks/enclave/gemma/colab-3-persona/2. DO-benchmark-owner-gemma-restrict.ipynb b/notebooks/enclave/gemma/colab-3-persona/2. DO-benchmark-owner-gemma-restrict.ipynb index 94d51734ef7..58dbe0f80ac 100644 --- a/notebooks/enclave/gemma/colab-3-persona/2. DO-benchmark-owner-gemma-restrict.ipynb +++ b/notebooks/enclave/gemma/colab-3-persona/2. DO-benchmark-owner-gemma-restrict.ipynb @@ -138,7 +138,7 @@ "source": [ "### Step 1.1 — Wait for the Researcher peer request, then approve\n", "\n", - "The Researcher notebook adds you as a peer. Re-run the cell below until you see their request appear, then approve." + "The Researcher notebook adds you as a peer. The cell below shows your peers. `peer_must_exist=False` lets you approve before their request arrives, and the next cell waits until the connection is live." ] }, { @@ -171,11 +171,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "benchmark_owner.validate_peer(RESEARCHER_EMAIL)\n", - "print(f\" Connected to {RESEARCHER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [RESEARCHER_EMAIL, ENCLAVE_EMAIL]:\n", + " benchmark_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -195,8 +195,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", - "benchmark_owner.attest_peer(ENCLAVE_EMAIL)" + "# Wait until the enclave publishes its attestation, then verify it.\n", + "benchmark_owner.wait_until_attested(ENCLAVE_EMAIL)" ] }, { @@ -310,8 +310,7 @@ "outputs": [], "source": [ "RESTRICT_JOB_NAME = \"restrict_engine_review\"\n", - "benchmark_owner.sync()\n", - "restrict_job = next(j for j in benchmark_owner.jobs if j.name == RESTRICT_JOB_NAME)\n", + "restrict_job = benchmark_owner.wait_until_has_job(RESTRICT_JOB_NAME, where=lambda j: j.can_approve)\n", "print(f\" Benchmark owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" ] }, @@ -354,8 +353,9 @@ "metadata": {}, "outputs": [], "source": [ - "benchmark_owner.sync()\n", - "restrict_job = next(j for j in benchmark_owner.jobs if j.name == RESTRICT_JOB_NAME)\n", + "restrict_job = benchmark_owner.wait_until_has_job(\n", + " RESTRICT_JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", "outputs = {p.name: p for p in restrict_job.output_paths}\n", "print(f\" Output files : {list(outputs)}\")\n", "\n", @@ -383,8 +383,7 @@ "outputs": [], "source": [ "JOB_NAME = \"safety_eval_job\"\n", - "benchmark_owner.sync()\n", - "benchmark_owner_job = next(j for j in benchmark_owner.jobs if j.name == JOB_NAME)\n", + "benchmark_owner_job = benchmark_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve, timeout=900)\n", "print(f\" Benchmark owner sees '{JOB_NAME}' status={benchmark_owner_job.status}\")" ] }, diff --git a/notebooks/enclave/gemma/colab-3-persona/3. DS-researcher-gemma-restrict.ipynb b/notebooks/enclave/gemma/colab-3-persona/3. DS-researcher-gemma-restrict.ipynb index 30505d2a4ab..c522f7a6a58 100644 --- a/notebooks/enclave/gemma/colab-3-persona/3. DS-researcher-gemma-restrict.ipynb +++ b/notebooks/enclave/gemma/colab-3-persona/3. DS-researcher-gemma-restrict.ipynb @@ -137,7 +137,7 @@ "source": [ "### Step 1.1 — Wait until both Data Owners approve, then verify peers\n", "\n", - "Each DO needs to run their `approve_peer_request` cell. Re-run the cell below until both appear as approved peers." + "Each DO needs to run their `approve_peer_request` cell. The cell below shows your peers, and the next cell waits until both have approved." ] }, { @@ -158,11 +158,10 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "for peer_email in [MODEL_OWNER_EMAIL, BENCHMARK_OWNER_EMAIL]:\n", - " researcher.validate_peer(peer_email)\n", + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [MODEL_OWNER_EMAIL, BENCHMARK_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " researcher.wait_until_peered(peer_email)\n", " print(f\" Connected to {peer_email}\")" ] }, @@ -181,8 +180,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", - "researcher.attest_peer(ENCLAVE_EMAIL)" + "# Wait until the enclave publishes its attestation, then verify it.\n", + "researcher.wait_until_attested(ENCLAVE_EMAIL)" ] }, { @@ -200,7 +199,8 @@ "metadata": {}, "outputs": [], "source": [ - "researcher.sync()\n", + "researcher.wait_until_has_dataset(\"gemma3_model\", datasite=MODEL_OWNER_EMAIL)\n", + "researcher.wait_until_has_dataset(\"safety_prompts\", datasite=BENCHMARK_OWNER_EMAIL)\n", "researcher.datasets" ] }, @@ -341,7 +341,7 @@ "source": [ "## Step 6 — Wait for both DOs to approve and the enclave to run\n", "\n", - "Each DO must approve from their notebook; once both approve, the enclave runs the job and pushes results back. Re-sync until status is `\"done\"`." + "Each DO must approve from their notebook; once both approve, the enclave runs the job and pushes results back. The cell below waits until the status is `\"done\"`." ] }, { @@ -351,13 +351,11 @@ "metadata": {}, "outputs": [], "source": [ - "researcher.sync()\n", - "researcher_job = next(j for j in researcher.jobs if j.name == JOB_NAME)\n", + "researcher_job = researcher.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", "print(f\" Researcher job status : {researcher_job.status}\")\n", - "print(f\" Output files : {researcher_job.output_paths}\")\n", - "\n", - "assert researcher_job.status == \"done\", researcher_job.status\n", - "assert len(researcher_job.output_paths) > 0" + "print(f\" Output files : {researcher_job.output_paths}\")" ] }, { diff --git a/notebooks/enclave/gemma/colab-shield/1. DO-model-owner-gemma-shield.ipynb b/notebooks/enclave/gemma/colab-shield/1. DO-model-owner-gemma-shield.ipynb index 9819051ae03..d5940be0c86 100644 --- a/notebooks/enclave/gemma/colab-shield/1. DO-model-owner-gemma-shield.ipynb +++ b/notebooks/enclave/gemma/colab-shield/1. DO-model-owner-gemma-shield.ipynb @@ -218,7 +218,7 @@ "source": [ "### Step 1.1 — Wait for the Benchmark Owner peer request, then approve\n", "\n", - "The Benchmark Owner notebook adds you as a peer. Re-run the cell below until you see their request appear, then approve." + "The Benchmark Owner notebook adds you as a peer. The cell below shows your peers. `peer_must_exist=False` lets you approve before their request arrives, and the next cell waits until the connection is live." ] }, { @@ -249,11 +249,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "model_owner.validate_peer(BENCHMARK_OWNER_EMAIL)\n", - "print(f\" Connected to {BENCHMARK_OWNER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [BENCHMARK_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " model_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -273,8 +273,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", - "model_owner.attest_peer(ENCLAVE_EMAIL)" + "# Wait until the enclave publishes its attestation, then verify it.\n", + "model_owner.wait_until_attested(ENCLAVE_EMAIL)" ] }, { @@ -645,12 +645,8 @@ "metadata": {}, "outputs": [], "source": [ - "restrict_job = next((j for j in model_owner.jobs if j.name == RESTRICT_JOB_NAME), None)\n", - "\n", - "if restrict_job is None:\n", - " print(f\" 🟠 Job '{RESTRICT_JOB_NAME}' not visible yet — the enclave is still distributing it. Wait a moment and re-run this cell.\")\n", - "else:\n", - " print(f\" ✅ Model owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" + "restrict_job = model_owner.wait_until_has_job(RESTRICT_JOB_NAME, where=lambda j: j.can_approve)\n", + "print(f\" ✅ Model owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" ] }, { @@ -700,12 +696,8 @@ "outputs": [], "source": [ "JOB_NAME = \"safety_eval_job\"\n", - "model_owner_job = next((j for j in model_owner.jobs if j.name == JOB_NAME), None)\n", - "\n", - "if model_owner_job is None:\n", - " print(f\" 🟠 Job '{JOB_NAME}' not visible yet — the enclave is still distributing it. Wait a moment and re-run this cell.\")\n", - "else:\n", - " print(f\" ✅ Model owner sees '{JOB_NAME}' status={model_owner_job.status}\")" + "model_owner_job = model_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve, timeout=900)\n", + "print(f\" ✅ Model owner sees '{JOB_NAME}' status={model_owner_job.status}\")" ] }, { diff --git a/notebooks/enclave/gemma/colab-shield/2. DO-benchmark-owner-gemma-shield.ipynb b/notebooks/enclave/gemma/colab-shield/2. DO-benchmark-owner-gemma-shield.ipynb index c6f01bdf378..04d449d4f66 100644 --- a/notebooks/enclave/gemma/colab-shield/2. DO-benchmark-owner-gemma-shield.ipynb +++ b/notebooks/enclave/gemma/colab-shield/2. DO-benchmark-owner-gemma-shield.ipynb @@ -171,7 +171,7 @@ "source": [ "### Step 1.1 — Wait until the Model Owner approves\n", "\n", - "The Model Owner approves your request from their notebook. Re-run the cell below until they appear as an approved peer." + "The Model Owner approves your request from their notebook. The cell below shows your peers, and the next cell waits until the Model Owner has approved." ] }, { @@ -191,11 +191,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "benchmark_owner.validate_peer(MODEL_OWNER_EMAIL)\n", - "print(f\" Connected to {MODEL_OWNER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [MODEL_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " benchmark_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -215,8 +215,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", - "benchmark_owner.attest_peer(ENCLAVE_EMAIL)" + "# Wait until the enclave publishes its attestation, then verify it.\n", + "benchmark_owner.wait_until_attested(ENCLAVE_EMAIL)" ] }, { @@ -411,12 +411,8 @@ "outputs": [], "source": [ "RESTRICT_JOB_NAME = \"restrict_engine_review\"\n", - "restrict_job = next((j for j in benchmark_owner.jobs if j.name == RESTRICT_JOB_NAME), None)\n", - "\n", - "if restrict_job is None:\n", - " print(f\" 🟠 Job '{RESTRICT_JOB_NAME}' not visible yet — the enclave is still distributing it. Wait a moment and re-run this cell.\")\n", - "else:\n", - " print(f\" Benchmark owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" + "restrict_job = benchmark_owner.wait_until_has_job(RESTRICT_JOB_NAME, where=lambda j: j.can_approve)\n", + "print(f\" Benchmark owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" ] }, { @@ -457,18 +453,17 @@ "metadata": {}, "outputs": [], "source": [ - "restrict_job = next(j for j in benchmark_owner.jobs if j.name == RESTRICT_JOB_NAME)\n", + "restrict_job = benchmark_owner.wait_until_has_job(\n", + " RESTRICT_JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", "\n", - "if restrict_job.status != \"done\":\n", - " print(f\" 🟠 Job '{RESTRICT_JOB_NAME}' status={restrict_job.status} — wait until it is 'done', then re-run this cell.\")\n", - "else:\n", - " outputs = {p.name: p for p in restrict_job.output_paths}\n", - " print(f\" Output files : {list(outputs)}\")\n", + "outputs = {p.name: p for p in restrict_job.output_paths}\n", + "print(f\" Output files : {list(outputs)}\")\n", "\n", - " cert = json.loads(outputs[\"gemma_inference.certificate.json\"].read_text())\n", - " print(f\" calls checked : {cert['n_calls_checked']}\")\n", - " print(f\" policy id : {cert['policy_id']}\")\n", - " print(f\" source sha256 : {cert['source_sha256']}\")" + "cert = json.loads(outputs[\"gemma_inference.certificate.json\"].read_text())\n", + "print(f\" calls checked : {cert['n_calls_checked']}\")\n", + "print(f\" policy id : {cert['policy_id']}\")\n", + "print(f\" source sha256 : {cert['source_sha256']}\")" ] }, { @@ -494,6 +489,8 @@ "metadata": {}, "outputs": [], "source": [ + "benchmark_owner.wait_until_has_dataset(\"gemma3_model\", datasite=MODEL_OWNER_EMAIL)\n", + "benchmark_owner.wait_until_has_dataset(\"shieldgemma_model\", datasite=MODEL_OWNER_EMAIL)\n", "benchmark_owner.datasets" ] }, @@ -703,8 +700,7 @@ "metadata": {}, "outputs": [], "source": [ - "benchmark_owner_job = next((j for j in benchmark_owner.jobs if j.name == JOB_NAME), None)\n", - "\n", + "benchmark_owner_job = benchmark_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve)\n", "print(f\" Benchmark owner sees '{JOB_NAME}' status={benchmark_owner_job.status}\")" ] }, @@ -744,7 +740,7 @@ "source": [ "## Step 7 — View the results\n", "\n", - "The results come back to us because we submitted the job. Re-sync until the status is `done`.\n", + "The results come back to us because we submitted the job. The cell below waits until the status is `done`.\n", "\n", "Declined prompts carry the declined message and no timing stats. Prompts that passed carry the output-gate score too, and either the Gemma 3 completion or, if the response violated the policy, the output-filtered notice in its place. The withheld text never leaves the enclave." ] @@ -756,13 +752,10 @@ "metadata": {}, "outputs": [], "source": [ - "benchmark_owner_job = next(j for j in benchmark_owner.jobs if j.name == JOB_NAME)\n", - "print(f\" Job status : {benchmark_owner_job.status}\")\n", - "\n", - "if benchmark_owner_job.status != \"done\" or not benchmark_owner_job.output_paths:\n", - " print(f\" 🟠 Job '{JOB_NAME}' is not finished yet — wait until the status is 'done', then re-run this cell.\")\n", - "else:\n", - " print(f\" ✅ Output files : {benchmark_owner_job.output_paths}\")" + "benchmark_owner_job = benchmark_owner.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", + "print(f\" ✅ Output files : {benchmark_owner_job.output_paths}\")" ] }, { diff --git a/notebooks/enclave/gemma/colab/1. DO-model-owner-gemma-restrict.ipynb b/notebooks/enclave/gemma/colab/1. DO-model-owner-gemma-restrict.ipynb index 57a995e8551..7172f738261 100644 --- a/notebooks/enclave/gemma/colab/1. DO-model-owner-gemma-restrict.ipynb +++ b/notebooks/enclave/gemma/colab/1. DO-model-owner-gemma-restrict.ipynb @@ -170,7 +170,7 @@ "source": [ "### Step 1.1 — Wait for the Benchmark Owner peer request, then approve\n", "\n", - "The Benchmark Owner notebook adds you as a peer. Re-run the cell below until you see their request appear, then approve." + "The Benchmark Owner notebook adds you as a peer. The cell below shows your peers. `peer_must_exist=False` lets you approve before their request arrives, and the next cell waits until the connection is live." ] }, { @@ -201,11 +201,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "model_owner.validate_peer(BENCHMARK_OWNER_EMAIL)\n", - "print(f\" Connected to {BENCHMARK_OWNER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [BENCHMARK_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " model_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -225,8 +225,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", - "model_owner.attest_peer(ENCLAVE_EMAIL)" + "# Wait until the enclave publishes its attestation, then verify it.\n", + "model_owner.wait_until_attested(ENCLAVE_EMAIL)" ] }, { @@ -504,12 +504,8 @@ "metadata": {}, "outputs": [], "source": [ - "restrict_job = next((j for j in model_owner.jobs if j.name == RESTRICT_JOB_NAME), None)\n", - "\n", - "if restrict_job is None:\n", - " print(f\" 🟠 Job '{RESTRICT_JOB_NAME}' not visible yet — the enclave is still distributing it. Wait a moment and re-run this cell.\")\n", - "else:\n", - " print(f\" ✅ Model owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" + "restrict_job = model_owner.wait_until_has_job(RESTRICT_JOB_NAME, where=lambda j: j.can_approve)\n", + "print(f\" ✅ Model owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" ] }, { @@ -559,12 +555,8 @@ "outputs": [], "source": [ "JOB_NAME = \"safety_eval_job\"\n", - "model_owner_job = next((j for j in model_owner.jobs if j.name == JOB_NAME), None)\n", - "\n", - "if model_owner_job is None:\n", - " print(f\" 🟠 Job '{JOB_NAME}' not visible yet — the enclave is still distributing it. Wait a moment and re-run this cell.\")\n", - "else:\n", - " print(f\" ✅ Model owner sees '{JOB_NAME}' status={model_owner_job.status}\")" + "model_owner_job = model_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve, timeout=900)\n", + "print(f\" ✅ Model owner sees '{JOB_NAME}' status={model_owner_job.status}\")" ] }, { diff --git a/notebooks/enclave/gemma/colab/2. DO-benchmark-owner-gemma-restrict.ipynb b/notebooks/enclave/gemma/colab/2. DO-benchmark-owner-gemma-restrict.ipynb index 25a1c9f0737..6ab0d7f817a 100644 --- a/notebooks/enclave/gemma/colab/2. DO-benchmark-owner-gemma-restrict.ipynb +++ b/notebooks/enclave/gemma/colab/2. DO-benchmark-owner-gemma-restrict.ipynb @@ -160,7 +160,7 @@ "source": [ "### Step 1.1 — Wait until the Model Owner approves\n", "\n", - "The Model Owner approves your request from their notebook. Re-run the cell below until they appear as an approved peer." + "The Model Owner approves your request from their notebook. The cell below shows your peers, and the next cell waits until the Model Owner has approved." ] }, { @@ -180,11 +180,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "benchmark_owner.validate_peer(MODEL_OWNER_EMAIL)\n", - "print(f\" Connected to {MODEL_OWNER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [MODEL_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " benchmark_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -204,8 +204,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", - "benchmark_owner.attest_peer(ENCLAVE_EMAIL)" + "# Wait until the enclave publishes its attestation, then verify it.\n", + "benchmark_owner.wait_until_attested(ENCLAVE_EMAIL)" ] }, { @@ -400,12 +400,8 @@ "outputs": [], "source": [ "RESTRICT_JOB_NAME = \"restrict_engine_review\"\n", - "restrict_job = next((j for j in benchmark_owner.jobs if j.name == RESTRICT_JOB_NAME), None)\n", - "\n", - "if restrict_job is None:\n", - " print(f\" 🟠 Job '{RESTRICT_JOB_NAME}' not visible yet — the enclave is still distributing it. Wait a moment and re-run this cell.\")\n", - "else:\n", - " print(f\" Benchmark owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" + "restrict_job = benchmark_owner.wait_until_has_job(RESTRICT_JOB_NAME, where=lambda j: j.can_approve)\n", + "print(f\" Benchmark owner sees '{RESTRICT_JOB_NAME}' status={restrict_job.status}\")" ] }, { @@ -446,18 +442,17 @@ "metadata": {}, "outputs": [], "source": [ - "restrict_job = next(j for j in benchmark_owner.jobs if j.name == RESTRICT_JOB_NAME)\n", + "restrict_job = benchmark_owner.wait_until_has_job(\n", + " RESTRICT_JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", "\n", - "if restrict_job.status != \"done\":\n", - " print(f\" 🟠 Job '{RESTRICT_JOB_NAME}' status={restrict_job.status} — wait until it is 'done', then re-run this cell.\")\n", - "else:\n", - " outputs = {p.name: p for p in restrict_job.output_paths}\n", - " print(f\" Output files : {list(outputs)}\")\n", + "outputs = {p.name: p for p in restrict_job.output_paths}\n", + "print(f\" Output files : {list(outputs)}\")\n", "\n", - " cert = json.loads(outputs[\"gemma_inference.certificate.json\"].read_text())\n", - " print(f\" calls checked : {cert['n_calls_checked']}\")\n", - " print(f\" policy id : {cert['policy_id']}\")\n", - " print(f\" source sha256 : {cert['source_sha256']}\")" + "cert = json.loads(outputs[\"gemma_inference.certificate.json\"].read_text())\n", + "print(f\" calls checked : {cert['n_calls_checked']}\")\n", + "print(f\" policy id : {cert['policy_id']}\")\n", + "print(f\" source sha256 : {cert['source_sha256']}\")" ] }, { @@ -480,6 +475,7 @@ "metadata": {}, "outputs": [], "source": [ + "benchmark_owner.wait_until_has_dataset(\"gemma3_model\", datasite=MODEL_OWNER_EMAIL)\n", "benchmark_owner.datasets" ] }, @@ -604,8 +600,7 @@ "metadata": {}, "outputs": [], "source": [ - "benchmark_owner_job = next((j for j in benchmark_owner.jobs if j.name == JOB_NAME), None)\n", - "\n", + "benchmark_owner_job = benchmark_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve)\n", "print(f\" Benchmark owner sees '{JOB_NAME}' status={benchmark_owner_job.status}\")" ] }, @@ -645,7 +640,7 @@ "source": [ "## Step 7 — View the results\n", "\n", - "The results come back to us because we submitted the job. Re-sync until the status is `done`." + "The results come back to us because we submitted the job. The cell below waits until the status is `done`." ] }, { @@ -655,13 +650,10 @@ "metadata": {}, "outputs": [], "source": [ - "benchmark_owner_job = next(j for j in benchmark_owner.jobs if j.name == JOB_NAME)\n", - "print(f\" Job status : {benchmark_owner_job.status}\")\n", - "\n", - "if benchmark_owner_job.status != \"done\" or not benchmark_owner_job.output_paths:\n", - " print(f\" 🟠 Job '{JOB_NAME}' is not finished yet — wait until the status is 'done', then re-run this cell.\")\n", - "else:\n", - " print(f\" ✅ Output files : {benchmark_owner_job.output_paths}\")" + "benchmark_owner_job = benchmark_owner.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", + "print(f\" ✅ Output files : {benchmark_owner_job.output_paths}\")" ] }, { diff --git a/notebooks/enclave/nanolm/2. enclave_eval_localenc.ipynb b/notebooks/enclave/nanolm/2. enclave_eval_localenc.ipynb index ed211eca3ab..054860e96dc 100644 --- a/notebooks/enclave/nanolm/2. enclave_eval_localenc.ipynb +++ b/notebooks/enclave/nanolm/2. enclave_eval_localenc.ipynb @@ -352,13 +352,17 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# Nothing in this notebook checks the enclave address.\n", - "researcher.validate_peer(MODEL_OWNER_EMAIL)\n", - "researcher.validate_peer(BENCHMARK_OWNER_EMAIL)\n", - "model_owner.validate_peer(RESEARCHER_EMAIL)\n", - "benchmark_owner.validate_peer(RESEARCHER_EMAIL)\n", + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "# Each call waits at most 300 s, so the worst case for this cell is\n", + "# 35 minutes (7 x 300 s).\n", + "for client, peer_emails in [\n", + " (researcher, [MODEL_OWNER_EMAIL, BENCHMARK_OWNER_EMAIL, settings.email]),\n", + " (model_owner, [RESEARCHER_EMAIL, settings.email]),\n", + " (benchmark_owner, [RESEARCHER_EMAIL, settings.email]),\n", + "]:\n", + " for peer_email in peer_emails:\n", + " client.wait_until_peered(peer_email)\n", "print(\" Peers validated\")" ] }, @@ -438,7 +442,10 @@ "metadata": {}, "outputs": [], "source": [ - "researcher.sync()\n", + "# The data owners uploaded both datasets in this notebook, so timeout=0 checks\n", + "# once, and fails at once if a dataset is not there.\n", + "researcher.wait_until_has_dataset(\"gemma_model\", datasite=MODEL_OWNER_EMAIL, timeout=0)\n", + "researcher.wait_until_has_dataset(\"eval_benchmark\", datasite=BENCHMARK_OWNER_EMAIL, timeout=0)\n", "researcher.datasets" ] }, @@ -550,11 +557,11 @@ "source": [ "runner.tick()\n", "\n", - "model_owner.sync()\n", - "benchmark_owner.sync()\n", + "# The enclave runs in this notebook, so only runner.tick() moves the job.\n", + "# timeout=0 checks once, and fails at once if the job is not there.\n", "JOB_NAME = \"bias_eval_job\"\n", - "model_owner_job = next(j for j in model_owner.jobs if j.name == JOB_NAME)\n", - "benchmark_owner_job = next(j for j in benchmark_owner.jobs if j.name == JOB_NAME)\n", + "model_owner_job = model_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve, timeout=0)\n", + "benchmark_owner_job = benchmark_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve, timeout=0)\n", "print(f\" Model owner sees '{JOB_NAME}' status={model_owner_job.status}\")\n", "print(f\" Benchmark owner sees '{JOB_NAME}' status={benchmark_owner_job.status}\")" ] @@ -602,13 +609,11 @@ "source": [ "runner.tick()\n", "\n", - "researcher.sync()\n", - "researcher_job = next(j for j in researcher.jobs if j.name == JOB_NAME)\n", + "researcher_job = researcher.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=0\n", + ")\n", "print(f\" Researcher job status : {researcher_job.status}\")\n", - "print(f\" Output files : {researcher_job.output_paths}\")\n", - "\n", - "assert researcher_job.status == \"done\", researcher_job.status\n", - "assert len(researcher_job.output_paths) > 0" + "print(f\" Output files : {researcher_job.output_paths}\")" ] }, { @@ -656,17 +661,16 @@ "metadata": {}, "outputs": [], "source": [ - "model_owner.sync()\n", - "benchmark_owner.sync()\n", + "def done_with_outputs(client):\n", + " return client.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=0\n", + " )\n", "\n", - "model_owner_job = next(j for j in model_owner.jobs if j.name == JOB_NAME)\n", - "benchmark_owner_job = next(j for j in benchmark_owner.jobs if j.name == JOB_NAME)\n", "\n", + "model_owner_job = done_with_outputs(model_owner)\n", + "benchmark_owner_job = done_with_outputs(benchmark_owner)\n", "print(f\" Model owner — output files : {model_owner_job.output_paths}\")\n", - "print(f\" Benchmark owner — output files : {benchmark_owner_job.output_paths}\")\n", - "\n", - "assert len(model_owner_job.output_paths) > 0\n", - "assert len(benchmark_owner_job.output_paths) > 0" + "print(f\" Benchmark owner — output files : {benchmark_owner_job.output_paths}\")" ] } ], diff --git a/notebooks/enclave/nanolm/3. enclave_eval_e2e.ipynb b/notebooks/enclave/nanolm/3. enclave_eval_e2e.ipynb index 2c3e521fc53..fc24f7c314d 100644 --- a/notebooks/enclave/nanolm/3. enclave_eval_e2e.ipynb +++ b/notebooks/enclave/nanolm/3. enclave_eval_e2e.ipynb @@ -224,13 +224,17 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# Nothing in this notebook checks the enclave address.\n", - "researcher.validate_peer(MODEL_OWNER_EMAIL)\n", - "researcher.validate_peer(BENCHMARK_OWNER_EMAIL)\n", - "model_owner.validate_peer(RESEARCHER_EMAIL)\n", - "benchmark_owner.validate_peer(RESEARCHER_EMAIL)\n", + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "# Each call waits at most 300 s, so the worst case for this cell is\n", + "# 35 minutes (7 x 300 s).\n", + "for client, peer_emails in [\n", + " (researcher, [MODEL_OWNER_EMAIL, BENCHMARK_OWNER_EMAIL, settings.email]),\n", + " (model_owner, [RESEARCHER_EMAIL, settings.email]),\n", + " (benchmark_owner, [RESEARCHER_EMAIL, settings.email]),\n", + "]:\n", + " for peer_email in peer_emails:\n", + " client.wait_until_peered(peer_email)\n", "print(\" Peers validated\")" ] }, @@ -421,7 +425,8 @@ "metadata": {}, "outputs": [], "source": [ - "researcher.sync()\n", + "researcher.wait_until_has_dataset(\"gemma_model\", datasite=MODEL_OWNER_EMAIL)\n", + "researcher.wait_until_has_dataset(\"eval_benchmark\", datasite=BENCHMARK_OWNER_EMAIL)\n", "researcher.datasets" ] }, @@ -531,11 +536,8 @@ "metadata": {}, "outputs": [], "source": [ - "\n", - "model_owner.sync()\n", - "benchmark_owner.sync()\n", - "model_owner_job = next(j for j in model_owner.jobs if j.name == JOB_NAME)\n", - "benchmark_owner_job = next(j for j in benchmark_owner.jobs if j.name == JOB_NAME)\n", + "model_owner_job = model_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve)\n", + "benchmark_owner_job = benchmark_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve)\n", "print(f\" Model owner sees '{JOB_NAME}' status={model_owner_job.status}\")\n", "print(f\" Benchmark owner sees '{JOB_NAME}' status={benchmark_owner_job.status}\")" ] @@ -581,13 +583,11 @@ "metadata": {}, "outputs": [], "source": [ - "researcher.sync()\n", - "researcher_job = next(j for j in researcher.jobs if j.name == JOB_NAME)\n", + "researcher_job = researcher.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", "print(f\" Researcher job status : {researcher_job.status}\")\n", - "print(f\" Output files : {researcher_job.output_paths}\")\n", - "\n", - "assert researcher_job.status == \"done\", researcher_job.status\n", - "assert len(researcher_job.output_paths) > 0" + "print(f\" Output files : {researcher_job.output_paths}\")" ] }, { @@ -635,17 +635,16 @@ "metadata": {}, "outputs": [], "source": [ - "model_owner.sync()\n", - "benchmark_owner.sync()\n", + "def done_with_outputs(client):\n", + " return client.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths)\n", + " )\n", "\n", - "model_owner_job = next(j for j in model_owner.jobs if j.name == JOB_NAME)\n", - "benchmark_owner_job = next(j for j in benchmark_owner.jobs if j.name == JOB_NAME)\n", "\n", + "model_owner_job = done_with_outputs(model_owner)\n", + "benchmark_owner_job = done_with_outputs(benchmark_owner)\n", "print(f\" Model owner — output files : {model_owner_job.output_paths}\")\n", - "print(f\" Benchmark owner — output files : {benchmark_owner_job.output_paths}\")\n", - "\n", - "assert len(model_owner_job.output_paths) > 0\n", - "assert len(benchmark_owner_job.output_paths) > 0" + "print(f\" Benchmark owner — output files : {benchmark_owner_job.output_paths}\")" ] } ], diff --git a/notebooks/enclave/nanolm/nbsplit/1. DO-model-owner-nanolm.ipynb b/notebooks/enclave/nanolm/nbsplit/1. DO-model-owner-nanolm.ipynb index d3042fc417e..d3e9923d1e2 100644 --- a/notebooks/enclave/nanolm/nbsplit/1. DO-model-owner-nanolm.ipynb +++ b/notebooks/enclave/nanolm/nbsplit/1. DO-model-owner-nanolm.ipynb @@ -133,7 +133,7 @@ "source": [ "### Step 1.1 — Wait for the Researcher peer request, then approve\n", "\n", - "The Researcher notebook adds you as a peer. Re-run the cell below until you see their request appear, then approve." + "The Researcher notebook adds you as a peer. The cell below shows your peers. `peer_must_exist=False` lets you approve before their request arrives, and the next cell waits until the connection is live." ] }, { @@ -166,11 +166,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "model_owner.validate_peer(RESEARCHER_EMAIL)\n", - "print(f\" Connected to {RESEARCHER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [RESEARCHER_EMAIL, ENCLAVE_EMAIL]:\n", + " model_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -188,7 +188,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", + "# Check the enclave's attestation, if it publishes one.\n", "model_owner.attest_peer(ENCLAVE_EMAIL)" ] }, @@ -304,7 +304,7 @@ "source": [ "## Step 3 — Wait for the Researcher to submit the job, then approve\n", "\n", - "The Researcher submits `bias_eval_job` to the enclave. Re-sync until it appears here, inspect it, then approve." + "The Researcher submits `bias_eval_job` to the enclave. The cell below waits until it appears here. Inspect it, then approve." ] }, { @@ -315,8 +315,7 @@ "outputs": [], "source": [ "JOB_NAME = \"bias_eval_job\"\n", - "model_owner.sync()\n", - "model_owner_job = next(j for j in model_owner.jobs if j.name == JOB_NAME)\n", + "model_owner_job = model_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve, timeout=900)\n", "print(f\" Model owner sees '{JOB_NAME}' status={model_owner_job.status}\")" ] }, @@ -349,10 +348,10 @@ "metadata": {}, "outputs": [], "source": [ - "model_owner.sync()\n", - "model_owner_job = next(j for j in model_owner.jobs if j.name == JOB_NAME)\n", - "print(f\" Output files : {model_owner_job.output_paths}\")\n", - "assert len(model_owner_job.output_paths) > 0" + "model_owner_job = model_owner.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", + "print(f\" Output files : {model_owner_job.output_paths}\")" ] } ], diff --git a/notebooks/enclave/nanolm/nbsplit/2. DO-benchmark-owner-nanolm.ipynb b/notebooks/enclave/nanolm/nbsplit/2. DO-benchmark-owner-nanolm.ipynb index 490d9841f6a..14e6987118f 100644 --- a/notebooks/enclave/nanolm/nbsplit/2. DO-benchmark-owner-nanolm.ipynb +++ b/notebooks/enclave/nanolm/nbsplit/2. DO-benchmark-owner-nanolm.ipynb @@ -133,7 +133,7 @@ "source": [ "### Step 1.1 — Wait for the Researcher peer request, then approve\n", "\n", - "The Researcher notebook adds you as a peer. Re-run the cell below until you see their request appear, then approve." + "The Researcher notebook adds you as a peer. The cell below shows your peers. `peer_must_exist=False` lets you approve before their request arrives, and the next cell waits until the connection is live." ] }, { @@ -166,11 +166,11 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "benchmark_owner.validate_peer(RESEARCHER_EMAIL)\n", - "print(f\" Connected to {RESEARCHER_EMAIL}\")" + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [RESEARCHER_EMAIL, ENCLAVE_EMAIL]:\n", + " benchmark_owner.wait_until_peered(peer_email)\n", + " print(f\" Connected to {peer_email}\")" ] }, { @@ -188,7 +188,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", + "# Check the enclave's attestation, if it publishes one.\n", "benchmark_owner.attest_peer(ENCLAVE_EMAIL)" ] }, @@ -283,7 +283,7 @@ "source": [ "## Step 3 — Wait for the Researcher to submit the job, then approve\n", "\n", - "The Researcher submits `bias_eval_job` to the enclave. Re-sync until it appears here, inspect it, then approve." + "The Researcher submits `bias_eval_job` to the enclave. The cell below waits until it appears here. Inspect it, then approve." ] }, { @@ -294,8 +294,7 @@ "outputs": [], "source": [ "JOB_NAME = \"bias_eval_job\"\n", - "benchmark_owner.sync()\n", - "benchmark_owner_job = next(j for j in benchmark_owner.jobs if j.name == JOB_NAME)\n", + "benchmark_owner_job = benchmark_owner.wait_until_has_job(JOB_NAME, where=lambda j: j.can_approve, timeout=900)\n", "print(f\" Benchmark owner sees '{JOB_NAME}' status={benchmark_owner_job.status}\")" ] }, @@ -328,10 +327,10 @@ "metadata": {}, "outputs": [], "source": [ - "benchmark_owner.sync()\n", - "benchmark_owner_job = next(j for j in benchmark_owner.jobs if j.name == JOB_NAME)\n", - "print(f\" Output files : {benchmark_owner_job.output_paths}\")\n", - "assert len(benchmark_owner_job.output_paths) > 0" + "benchmark_owner_job = benchmark_owner.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", + "print(f\" Output files : {benchmark_owner_job.output_paths}\")" ] } ], diff --git a/notebooks/enclave/nanolm/nbsplit/3. DS-researcher-nanolm.ipynb b/notebooks/enclave/nanolm/nbsplit/3. DS-researcher-nanolm.ipynb index 3314f5eec98..45caa9c9f5e 100644 --- a/notebooks/enclave/nanolm/nbsplit/3. DS-researcher-nanolm.ipynb +++ b/notebooks/enclave/nanolm/nbsplit/3. DS-researcher-nanolm.ipynb @@ -137,7 +137,7 @@ "source": [ "### Step 1.1 — Wait for the Data Owners to approve the peer request\n", "\n", - "The Model Owner and Benchmark Owner each need to run their `approve_peer_request` cell. Re-run the cell below until both appear as approved peers." + "The Model Owner and Benchmark Owner each need to run their `approve_peer_request` cell. The cell below shows your peers, and the next cell waits until both have approved." ] }, { @@ -158,11 +158,10 @@ "metadata": {}, "outputs": [], "source": [ - "# Fail on a wrong email, or if the peer has not approved yet. Warn if the peer\n", - "# was already there before this notebook logged in.\n", - "# attest_peer() below checks the enclave.\n", - "for peer_email in [MODEL_OWNER_EMAIL, BENCHMARK_OWNER_EMAIL]:\n", - " researcher.validate_peer(peer_email)\n", + "# Wait up to 5 minutes for each peer to approve. Fail at once on a wrong email.\n", + "# Warn if the peer was already there before this notebook logged in.\n", + "for peer_email in [MODEL_OWNER_EMAIL, BENCHMARK_OWNER_EMAIL, ENCLAVE_EMAIL]:\n", + " researcher.wait_until_peered(peer_email)\n", " print(f\" Connected to {peer_email}\")" ] }, @@ -181,7 +180,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Wait for enclave to accept peer request\n", + "# Check the enclave's attestation, if it publishes one.\n", "researcher.attest_peer(ENCLAVE_EMAIL)" ] }, @@ -200,7 +199,8 @@ "metadata": {}, "outputs": [], "source": [ - "researcher.sync()\n", + "researcher.wait_until_has_dataset(\"nanolm_model\", datasite=MODEL_OWNER_EMAIL)\n", + "researcher.wait_until_has_dataset(\"eval_benchmark\", datasite=BENCHMARK_OWNER_EMAIL)\n", "researcher.datasets" ] }, @@ -330,7 +330,7 @@ "source": [ "## Step 6 — Wait for both DOs to approve and the enclave to run\n", "\n", - "Each DO must approve from their notebook; once both approve, the enclave runs the job and pushes results back. Re-sync until status is `\"done\"`." + "Each DO must approve from their notebook; once both approve, the enclave runs the job and pushes results back. The cell below waits until the status is `\"done\"`." ] }, { @@ -340,13 +340,11 @@ "metadata": {}, "outputs": [], "source": [ - "researcher.sync()\n", - "researcher_job = next(j for j in researcher.jobs if j.name == JOB_NAME)\n", + "researcher_job = researcher.wait_until_has_job(\n", + " JOB_NAME, status=\"done\", where=lambda j: bool(j.output_paths), timeout=900\n", + ")\n", "print(f\" Researcher job status : {researcher_job.status}\")\n", - "print(f\" Output files : {researcher_job.output_paths}\")\n", - "\n", - "assert researcher_job.status == \"done\", researcher_job.status\n", - "assert len(researcher_job.output_paths) > 0" + "print(f\" Output files : {researcher_job.output_paths}\")" ] }, { diff --git a/packages/syft-enclave/src/syft_enclaves/client.py b/packages/syft-enclave/src/syft_enclaves/client.py index 7d428ed2f75..8cbbe109a66 100644 --- a/packages/syft-enclave/src/syft_enclaves/client.py +++ b/packages/syft-enclave/src/syft_enclaves/client.py @@ -5,6 +5,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Optional +from syft_datasets import Dataset from syft_datasets.dataset_manager import SyftDatasetManager from syft_job.disclosures import DisclosuresArg, gated_names from syft_job.job import JobInfo, JobsList @@ -12,9 +13,15 @@ from syft_job.models import JobState, JobStatus from syft_perms.syftperm_context import SyftPermContext from syft_rds import SyftRDSClient, SyftRDSClientConfig +from syft_rds.waiting import JobCheck, JobStatusArg, wait_for_job from syft.sync.peers.peer import Peer from syft.sync.peers.peer_list import PeerList +from syft.sync.utils.waiting import ( + DEFAULT_POLL_INTERVAL, + DEFAULT_WAIT_TIMEOUT, + wait_for, +) from syft.sync.version.peer_manager import CompatAction from syft_enclaves.attestation.dispatch import policy_for, verify_evidence from syft_enclaves.attestation.envelope import AttestationEvidence @@ -62,6 +69,28 @@ def pre_sync_enabled() -> bool: logger = logging.getLogger(__name__) +def _given_shorthands( + expected_image_digest: str | None, + expected_data_owners: list[str] | None, + expected_email: str | None, + policy: "AppraisalPolicy | TinfoilAppraisalPolicy | None", +) -> dict: + """The shorthand arguments of ``attest_peer`` that were given. + + Raises: + ValueError: a shorthand and ``policy`` were both given. + """ + shorthands = { + "expected_image_digest": expected_image_digest, + "expected_data_owners": expected_data_owners, + "expected_email": expected_email, + } + given = {name: value for name, value in shorthands.items() if value is not None} + if given and policy is not None: + raise ValueError(f"Pass either {' / '.join(shorthands)} or policy, not both.") + return given + + class SyftEnclaveClient: def __init__( self, @@ -151,17 +180,9 @@ def attest_peer( the shorthands. """ - shorthands = { - "expected_image_digest": expected_image_digest, - "expected_data_owners": expected_data_owners, - "expected_email": expected_email, - } - given = {name: value for name, value in shorthands.items() if value is not None} - if given and policy is not None: - raise ValueError( - f"Pass either {' / '.join(shorthands)} or policy, not both." - ) - + given = _given_shorthands( + expected_image_digest, expected_data_owners, expected_email, policy + ) evidence = self._peer_evidence(peer_email) if evidence is None: return None @@ -205,21 +226,65 @@ def _persist_peer_bundle(self, peer_email: str, bundle: dict) -> None: peer_email, peer.state.value, public_encryption_bundle=bundle ) - def _peer_evidence(self, peer_email: str) -> "AttestationEvidence | None": - """The peer's published attestation evidence, or None if it has none.""" + def wait_until_attested( + self, + peer_email: str, + expected_image_digest: str | None = None, + expected_data_owners: list[str] | None = None, + expected_email: str | None = None, + policy: "AppraisalPolicy | TinfoilAppraisalPolicy | None" = None, + timeout: float = DEFAULT_WAIT_TIMEOUT, + poll_interval: float = DEFAULT_POLL_INTERVAL, + ): + """Wait until ``peer_email`` publishes attestation evidence, then verify it. + + The arguments and the result are those of ``attest_peer``. Use it only + for an enclave that runs in a TEE: an enclave outside a TEE never + publishes evidence, so the wait ends in ``TimeoutError``. + + Raises: + TimeoutError: no evidence was published within ``timeout`` seconds. + AttestationError: the evidence fails verification. + """ + _given_shorthands( + expected_image_digest, expected_data_owners, expected_email, policy + ) + wait_for( + lambda: self._peer_evidence(peer_email, quiet=True), + f"the attestation of {peer_email}", + lambda: "no attestation evidence published yet", + timeout, + poll_interval, + ) + return self.attest_peer( + peer_email, + expected_image_digest=expected_image_digest, + expected_data_owners=expected_data_owners, + expected_email=expected_email, + policy=policy, + ) + + def _peer_evidence( + self, peer_email: str, quiet: bool = False + ) -> "AttestationEvidence | None": + """The peer's published attestation evidence, or None if it has none. + + Prints why there is none, unless ``quiet`` is set. + """ version_info = self._rds.peer_manager.connection_router.read_peer_version_file( peer_email ) if version_info is None: - print( - f"ℹī¸ No version file available for peer {peer_email!r}; skipping attestation." - ) + if not quiet: + print( + f"ℹī¸ No version file available for peer {peer_email!r}; skipping attestation." + ) return None # A malformed envelope raises rather than skipping: a peer that # published something unparseable is not the same as one that # published nothing. evidence = AttestationEvidence.read_from(version_info) - if evidence is None: + if evidence is None and not quiet: print( f"ℹī¸ Peer {peer_email!r} published no attestation evidence " "(not running in an attested enclave); skipping attestation." @@ -249,13 +314,64 @@ def datasets(self) -> SyftDatasetManager: @property def jobs(self) -> JobsList: - jobs_list = self._rds.jobs + return self._wrap_enclave_jobs(self._rds.jobs) + + def _wrap_enclave_jobs(self, jobs_list: JobsList) -> JobsList: + """``jobs_list`` with every enclave job wrapped by ``_as_enclave_job``.""" wrapped = [ self._as_enclave_job(j) if j.job_headers.get("job_type") == "enclave" else j for j in jobs_list ] return JobsList(wrapped, jobs_list._root_email) + def wait_until_peered( + self, + peer_email: str, + timeout: float = DEFAULT_WAIT_TIMEOUT, + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> Peer: + """Same as ``SyftRDSClient.wait_until_peered``.""" + return self._rds.wait_until_peered( + peer_email, timeout=timeout, poll_interval=poll_interval + ) + + def wait_until_has_dataset( + self, + name: str, + datasite: str | None = None, + timeout: float = DEFAULT_WAIT_TIMEOUT, + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> Dataset: + """Same as ``SyftRDSClient.wait_until_has_dataset``.""" + return self._rds.wait_until_has_dataset( + name, datasite=datasite, timeout=timeout, poll_interval=poll_interval + ) + + def wait_until_has_job( + self, + job_name: str, + user_name: str | None = None, + status: JobStatusArg = None, + where: JobCheck | None = None, + timeout: float = DEFAULT_WAIT_TIMEOUT, + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> JobInfo: + """Like ``SyftRDSClient.wait_until_has_job``, with enclave job statuses. + + An enclave job is read as in ``self.jobs``, so its status comes from + the approval files of the data owners. + """ + return wait_for_job( + self.sync, + lambda: self._wrap_enclave_jobs(self._local_jobs()), + job_name, + user_name, + status, + where, + timeout, + poll_interval, + ) + def _as_enclave_job(self, job: JobInfo) -> EnclaveJobInfo: """Wrap ``job``; on the enclave that runs it, every data owner must approve. diff --git a/packages/syft-enclave/tests/test_enclave_jobs.py b/packages/syft-enclave/tests/test_enclave_jobs.py index c9a707fb0d4..0ade9b35625 100644 --- a/packages/syft-enclave/tests/test_enclave_jobs.py +++ b/packages/syft-enclave/tests/test_enclave_jobs.py @@ -2,14 +2,18 @@ import os import random import tempfile +import threading +import time from pathlib import Path import pytest os.environ["PRE_SYNC"] = "false" +from syft.sync.connections.drive import mock_drive_service from syft_enclaves import SyftEnclaveClient from syft_enclaves.enclave_job_info import ( + EnclaveJobInfo, PartyApprovalStatus, enclave_approval_file_name, ) @@ -497,3 +501,126 @@ def test_rejection_rejects_job_and_withdraws_approval(): assert enclave.jobs["test_job"].status == "rejected" stored = PartyApprovalStatus.load_json(_approval_file(enclave, do2.email)) assert stored.reason == "uses more data than agreed" + + +def test_enclave_waits_until_every_data_owner_approves(monkeypatch): + """The wait reads the enclave's own job status, derived from the approval + files, not the raw job state, which stays pending.""" + enclave, do1, do2 = _job_distributed_to_both_data_owners() + sleeps = [] + + def approve_on_first_sleep(seconds): + sleeps.append(seconds) + if len(sleeps) == 1: + do1.approve_job(do1.jobs["test_job"]) + do2.approve_job(do2.jobs["test_job"]) + + monkeypatch.setattr(time, "sleep", approve_on_first_sleep) + + job = enclave.wait_until_has_job("test_job", status="approved") + assert isinstance(job, EnclaveJobInfo) + assert job.status == "approved" + assert enclave._local_jobs()["test_job"].status == "pending" + assert sleeps == [15] + + +def _serialize_mock_drive(monkeypatch): + """Run one mock Drive request at a time; the in-memory store is not thread-safe.""" + lock = threading.RLock() + for name in dir(mock_drive_service): + request_class = getattr(mock_drive_service, name) + if isinstance(request_class, type) and "execute" in vars(request_class): + + def execute(self, *args, _original=request_class.execute, **kwargs): + with lock: + return _original(self, *args, **kwargs) + + monkeypatch.setattr(request_class, "execute", execute) + + +def test_parties_run_concurrently_and_meet_through_waits(monkeypatch): + """Each party runs its cells in its own thread, as in a notebook "run all". + + Nothing orders the threads: only the wait_until_* helpers make each party + wait for the others. + """ + _serialize_mock_drive(monkeypatch) + enclave, do1, do2, ds = SyftEnclaveClient.quad_with_mock_drive_service_connection( + use_in_memory_cache=False, + ) + wait = {"timeout": 120, "poll_interval": 0.05} + done = threading.Event() + errors = [] + results = {} + + def data_owner(do, name, prefix): + mock, private = create_tmp_dataset_files(prefix) + do.create_dataset( + name=name, + mock_path=mock, + private_path=private, + summary=name, + users=[ds.email, enclave.email], + upload_private=True, + sync=False, + ) + do.share_private_dataset(name, enclave.email) + do.sync() + job = do.wait_until_has_job("test_job", where=lambda j: j.can_approve, **wait) + do.approve_job(job) + results[do.email] = do.wait_until_has_job( + "test_job", status="done", where=lambda j: bool(j.output_paths), **wait + ) + + def data_scientist(): + ds.wait_until_has_dataset("dataset1", datasite=do1.email, **wait) + ds.wait_until_has_dataset("dataset2", datasite=do2.email, **wait) + ds.submit_python_job( + enclave.email, + create_tmp_code_file(make_job_code(do1.email, do2.email)), + "test_job", + datasets={do1.email: ["dataset1"], do2.email: ["dataset2"]}, + share_results_with_do=True, + ) + results[ds.email] = ds.wait_until_has_job( + "test_job", status="done", where=lambda j: bool(j.output_paths), **wait + ) + + def run_enclave(): + while not done.is_set(): + enclave.sync() + enclave.receive_jobs() + enclave.run_jobs() + enclave.distribute_results() + done.wait(0.05) + + def guarded(target, *args): + def run(): + try: + target(*args) + except BaseException as e: # reported to the main thread below + errors.append(e) + done.set() + + return threading.Thread(target=run, daemon=True) + + enclave_thread = guarded(run_enclave) + parties = [ + guarded(data_owner, do1, "dataset1", "do1"), + guarded(data_owner, do2, "dataset2", "do2"), + guarded(data_scientist), + ] + enclave_thread.start() + for thread in parties: + thread.start() + for thread in parties: + thread.join(timeout=180) + done.set() + enclave_thread.join(timeout=30) + + if errors: + raise errors[0] + assert not any(t.is_alive() for t in parties) + for email in (ds.email, do1.email, do2.email): + assert results[email].status == "done" + assert results[email].output_paths diff --git a/packages/syft-enclave/tests/test_wait_until_attested.py b/packages/syft-enclave/tests/test_wait_until_attested.py new file mode 100644 index 00000000000..b665511a590 --- /dev/null +++ b/packages/syft-enclave/tests/test_wait_until_attested.py @@ -0,0 +1,93 @@ +import os +import time + +os.environ["PRE_SYNC"] = "false" + +import pytest + +from syft_enclaves import SyftEnclaveClient +from syft_enclaves.attestation.result import AttestationError + + +@pytest.fixture +def clock(monkeypatch): + """Fake monotonic clock; time.sleep() advances it and records each sleep.""" + state = {"now": 0.0, "sleeps": []} + + def sleep(seconds): + state["sleeps"].append(seconds) + state["now"] += seconds + + monkeypatch.setattr(time, "sleep", sleep) + monkeypatch.setattr(time, "monotonic", lambda: state["now"]) + return state + + +@pytest.fixture +def ds(): + enclave, _, _, ds = SyftEnclaveClient.quad_with_mock_drive_service_connection() + ds.enclave_email = enclave.email + return ds + + +def test_waits_for_evidence_then_verifies_it(ds, monkeypatch, clock): + published = iter([None, None, "evidence"]) + monkeypatch.setattr(ds, "_peer_evidence", lambda email, quiet: next(published)) + calls = [] + + def attest_peer(email, **kwargs): + calls.append((email, kwargs)) + return "verified" + + monkeypatch.setattr(ds, "attest_peer", attest_peer) + + result = ds.wait_until_attested(ds.enclave_email, policy="the policy") + assert result == "verified" + assert calls == [ + ( + ds.enclave_email, + { + "expected_image_digest": None, + "expected_data_owners": None, + "expected_email": None, + "policy": "the policy", + }, + ) + ] + assert clock["sleeps"] == [15, 15] + + +def test_times_out_quietly_when_no_evidence_arrives(ds, monkeypatch, clock, capsys): + monkeypatch.setattr(ds, "_peer_evidence", lambda email, quiet: None) + capsys.readouterr() + + with pytest.raises(TimeoutError, match="attestation"): + ds.wait_until_attested(ds.enclave_email, timeout=30) + assert "published no attestation evidence" not in capsys.readouterr().out + assert clock["sleeps"] == [15, 15] + + +def test_failed_verification_raises_at_once(ds, monkeypatch, clock): + monkeypatch.setattr(ds, "_peer_evidence", lambda email, quiet: "evidence") + + def attest_peer(email, **kwargs): + raise AttestationError("digest mismatch") + + monkeypatch.setattr(ds, "attest_peer", attest_peer) + + with pytest.raises(AttestationError, match="digest mismatch"): + ds.wait_until_attested(ds.enclave_email) + assert clock["sleeps"] == [] + + +def test_bad_arguments_fail_before_wait(ds, monkeypatch, clock): + def no_lookup(email, quiet): + raise AssertionError("must not look for evidence") + + monkeypatch.setattr(ds, "_peer_evidence", no_lookup) + + with pytest.raises(ValueError, match="not both"): + ds.wait_until_attested( + ds.enclave_email, expected_email=ds.enclave_email, policy="the policy" + ) + assert clock["sleeps"] == [] diff --git a/packages/syft-rds/src/syft_rds/client.py b/packages/syft-rds/src/syft_rds/client.py index 30983ae49bb..c30f4960f97 100644 --- a/packages/syft-rds/src/syft_rds/client.py +++ b/packages/syft-rds/src/syft_rds/client.py @@ -13,6 +13,8 @@ from pydantic import BaseModel, ConfigDict from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.peers.peer import Peer +from syft.sync.utils.waiting import DEFAULT_POLL_INTERVAL, DEFAULT_WAIT_TIMEOUT from syft_datasets.dataset_manager import ( DATASET_COLLECTION_PREFIX, PRIVATE_DATASET_COLLECTION_PREFIX, @@ -24,6 +26,7 @@ from syft_datasets.dataset_manager import SHARE_WITH_ANY, SyftDatasetManager from syft_datasets.dataset_ref import DatasetNotFoundError from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION +from syft_rds.waiting import JobCheck, JobStatusArg, wait_for_dataset, wait_for_job from syft_rds.apis import Api, ApiCollection from syft_rds.config import ( DATASET_COLLECTION_SPECS, @@ -496,6 +499,86 @@ def jobs(self) -> Any: self.sync_engine.sync() return self.job_client.jobs + # ------------------------------------------------------------------ # + # waits for another party (notebook "run all" without re-run cells) + # ------------------------------------------------------------------ # + def wait_until_peered( + self, + peer_email: str, + timeout: float = DEFAULT_WAIT_TIMEOUT, + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> Peer: + """Wait until ``peer_email`` is a live connection, as ``validate_peer`` checks. + + Raises: + TimeoutError: the peer did not approve within ``timeout`` seconds. + PeerSetupError: waiting cannot help, e.g. the email is not a peer. + """ + return self.sync_engine.wait_until_peered( + peer_email, timeout=timeout, poll_interval=poll_interval + ) + + def wait_until_has_dataset( + self, + name: str, + datasite: str | None = None, + timeout: float = DEFAULT_WAIT_TIMEOUT, + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> Any: + """Sync until the dataset ``name`` is here, and return it. + + ``datasite`` is the owner; it is needed only when more than one owner + has a dataset called ``name``. + + Raises: + TimeoutError: the dataset did not arrive within ``timeout`` seconds. + ValueError: more than one owner has a dataset called ``name``. + """ + return wait_for_dataset( + self.sync, + self.dataset_manager.get_all, + name, + datasite, + timeout, + poll_interval, + ) + + def wait_until_has_job( + self, + job_name: str, + user_name: str | None = None, + status: JobStatusArg = None, + where: JobCheck | None = None, + timeout: float = DEFAULT_WAIT_TIMEOUT, + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> Any: + """Sync until the job ``job_name`` is here with ``status``, and return it. + + ``user_name`` matches the submitter or the datasite owner; it is needed + only when more than one job has the name. ``status`` is one status or + several; None accepts any. A job that is done counts for any status + before it, because a poll can miss a short status. To wait for a + status before ``done``, also pass the statuses after it, for example + ``("approved", "running")``. ``where`` is an + extra condition on the job, for example + ``lambda job: bool(job.output_paths)``. + + Raises: + TimeoutError: the job did not match within ``timeout`` seconds. + JobEndedError: the job ended with a final status not in ``status``. + ValueError: more than one job matches, or ``status`` is not valid. + """ + return wait_for_job( + self.sync, + lambda: self.job_client.jobs, + job_name, + user_name, + status, + where, + timeout, + poll_interval, + ) + def process_approved_jobs( self, stream_output: bool = True, diff --git a/packages/syft-rds/src/syft_rds/waiting.py b/packages/syft-rds/src/syft_rds/waiting.py new file mode 100644 index 00000000000..509ee7029b5 --- /dev/null +++ b/packages/syft-rds/src/syft_rds/waiting.py @@ -0,0 +1,225 @@ +"""Find a job or a dataset by name, and wait until it is there. + +The waits sync the client, then read the local job or dataset list, at each +poll. The loop is ``syft.sync.utils.waiting.wait_for``. +""" + +from collections.abc import Callable, Iterable +from typing import Any, Optional + +from syft.sync.utils.waiting import wait_for +from syft_job.models import JobStatus + +# A job with one of these statuses does not change status again. +FINAL_JOB_STATUSES = frozenset( + {JobStatus.DONE.value, JobStatus.FAILED.value, JobStatus.REJECTED.value} +) + +JobStatusArg = str | JobStatus | Iterable[str | JobStatus] | None + +# An extra condition on a job, for example ``lambda job: job.can_approve``. +JobCheck = Callable[[Any], bool] + + +class JobEndedError(RuntimeError): + """The job reached a final status that the caller did not wait for.""" + + +def job_statuses(status: JobStatusArg) -> Optional[frozenset[str]]: + """The status values in ``status``, or None for "any status". + + Raises: + ValueError: a value is not a ``JobStatus``. + """ + if status is None: + return None + items = [status] if isinstance(status, (str, JobStatus)) else list(status) + return frozenset(JobStatus(item).value for item in items) + + +def _jobs_named(jobs: Iterable[Any], job_name: str, user_name: Optional[str]) -> list: + """Jobs called ``job_name``; with ``user_name``, only the ones it submitted or hosts.""" + return [ + job + for job in jobs + if job.name == job_name + and ( + user_name is None + or user_name in (job.submitted_by, job.datasite_owner_email) + ) + ] + + +def _has_status(job: Any, statuses: Optional[frozenset[str]]) -> bool: + """True when ``job`` has one of ``statuses``, or is done after one of them. + + Every status that is not final comes before ``done``. A job that is done + therefore passed any of them, also if no poll saw it. + """ + if statuses is None or job.status in statuses: + return True + return job.status == JobStatus.DONE.value and bool(statuses - FINAL_JOB_STATUSES) + + +def _single_job( + jobs: Iterable[Any], job_name: str, user_name: Optional[str] +) -> Optional[Any]: + """The job called ``job_name``, or None. + + Jobs with the name, the submitter and the datasite in common are runs of + the same job, for example from an earlier run of a notebook. The newest + run is the one returned. + + Raises: + ValueError: jobs with the name have more than one submitter or datasite. + """ + matches = _jobs_named(jobs, job_name, user_name) + origins = sorted({(job.submitted_by, job.datasite_owner_email) for job in matches}) + if len(origins) > 1: + candidates = ", ".join( + f"submitted by {submitter} to {datasite}" for submitter, datasite in origins + ) + raise ValueError( + f"{len(matches)} jobs are named {job_name!r} ({candidates}). " + "Pass user_name= to pick one." + ) + # submitted_at is an ISO 8601 string, so the newest sorts last. + return max(matches, key=lambda job: job.submitted_at or "", default=None) + + +def find_job( + jobs: Iterable[Any], + job_name: str, + user_name: Optional[str], + statuses: Optional[frozenset[str]], + where: Optional[JobCheck] = None, +) -> Optional[Any]: + """Return the job called ``job_name`` once it has one of ``statuses``. + + ``user_name`` matches the submitter or the datasite owner of the job. It is + needed only when more than one job has the name. With ``where``, the job + must also make ``where(job)`` True. Returns None while no job matches, or + while the job can still change. + + Raises: + ValueError: more than one job matches. + JobEndedError: the job has a final status that is not in ``statuses``. + """ + job = _single_job(jobs, job_name, user_name) + if job is None: + return None + has_status = _has_status(job, statuses) + if has_status and (where is None or where(job)): + return job + if not has_status and job.status in FINAL_JOB_STATUSES: + raise JobEndedError( + f"Job {job_name!r} ended as {job.status!r}, not " + f"{' or '.join(sorted(statuses))}." + ) + return None + + +def describe_job( + jobs: Iterable[Any], + job_name: str, + user_name: Optional[str], + statuses: Optional[frozenset[str]] = None, + where: Optional[JobCheck] = None, +) -> str: + """Say what the local job list holds for ``job_name`` now.""" + job = _single_job(jobs, job_name, user_name) + if job is None: + for_user = f" for {user_name}" if user_name else "" + return f"no job named {job_name!r}{for_user} yet" + text = f"job {job_name!r} is {job.status!r}" + has_status = _has_status(job, statuses) + if where is not None and has_status and not where(job): + text += ", and the where= check is False" + return text + + +def find_dataset( + datasets: Iterable[Any], name: str, datasite: Optional[str] +) -> Optional[Any]: + """Return the dataset called ``name``, owned by ``datasite`` when given. + + Raises: + ValueError: more than one owner has a dataset called ``name``. + """ + matches = [ + dataset + for dataset in datasets + if dataset.name == name and (datasite is None or dataset.owner == datasite) + ] + if not matches: + return None + if len(matches) > 1: + owners = ", ".join(dataset.owner for dataset in matches) + raise ValueError( + f"{len(matches)} datasets are named {name!r} (owned by {owners}). " + "Pass datasite= to pick one." + ) + return matches[0] + + +def wait_for_job( + sync: Callable[[], Any], + list_jobs: Callable[[], Iterable[Any]], + job_name: str, + user_name: Optional[str], + status: JobStatusArg, + where: Optional[JobCheck], + timeout: float, + poll_interval: float, +) -> Any: + """Sync, then look for the job, until ``find_job`` returns it. + + Raises: + TimeoutError: the job did not match within ``timeout`` seconds. + ValueError: more than one job matches, or ``status`` is not valid. + JobEndedError: the job has a final status that is not in ``status``. + """ + statuses = job_statuses(status) + wanted = f" to be {' or '.join(sorted(statuses))}" if statuses else "" + checked = " and to pass the where= check" if where else "" + + def find() -> Optional[Any]: + sync() + return find_job(list_jobs(), job_name, user_name, statuses, where) + + return wait_for( + find, + f"job {job_name!r}{wanted}{checked}", + lambda: describe_job(list_jobs(), job_name, user_name, statuses, where), + timeout, + poll_interval, + ) + + +def wait_for_dataset( + sync: Callable[[], Any], + list_datasets: Callable[[], Iterable[Any]], + name: str, + datasite: Optional[str], + timeout: float, + poll_interval: float, +) -> Any: + """Sync, then look for the dataset, until ``find_dataset`` returns it. + + Raises: + TimeoutError: no dataset matched within ``timeout`` seconds. + ValueError: more than one owner has a dataset called ``name``. + """ + from_owner = f" from {datasite}" if datasite else "" + + def find() -> Optional[Any]: + sync() + return find_dataset(list_datasets(), name, datasite) + + return wait_for( + find, + f"dataset {name!r}{from_owner}", + lambda: f"no dataset named {name!r}{from_owner} yet", + timeout, + poll_interval, + ) diff --git a/packages/syft-rds/tests/test_wait_helpers.py b/packages/syft-rds/tests/test_wait_helpers.py new file mode 100644 index 00000000000..fd7b457c539 --- /dev/null +++ b/packages/syft-rds/tests/test_wait_helpers.py @@ -0,0 +1,224 @@ +"""wait_until_* helpers: job and dataset matching, and the waits over mock Drive.""" + +import time +from types import SimpleNamespace + +import pytest +from dataset_test_utils import create_tmp_dataset_files +from syft_rds import SyftRDSClient +from syft_rds.waiting import JobEndedError, describe_job, find_dataset, find_job + + +def _job( + name="j", + submitted_by="ds@x", + datasite="do@x", + status="pending", + submitted_at="2026-09-30T10:00:00+00:00", +): + return SimpleNamespace( + name=name, + submitted_by=submitted_by, + datasite_owner_email=datasite, + status=status, + submitted_at=submitted_at, + ) + + +def _dataset(name="d", owner="do@x"): + return SimpleNamespace(name=name, owner=owner) + + +@pytest.fixture +def clock(monkeypatch): + """Fake monotonic clock; each time.sleep() runs the queued callbacks first.""" + state = {"now": 0.0, "sleeps": [], "on_sleep": []} + + def sleep(seconds): + state["sleeps"].append(seconds) + state["now"] += seconds + if state["on_sleep"]: + state["on_sleep"].pop(0)() + + monkeypatch.setattr(time, "sleep", sleep) + monkeypatch.setattr(time, "monotonic", lambda: state["now"]) + return state + + +# ---------------------------------------------------------------- find_job --- + + +def test_find_job_returns_only_job_with_name(): + job = _job() + assert find_job([job, _job(name="other")], "j", None, None) is job + + +def test_find_job_returns_none_when_no_job_has_name(): + assert find_job([_job(name="other")], "j", None, None) is None + + +def test_find_job_raises_when_name_has_two_submitters(): + jobs = [_job(submitted_by="a@x"), _job(submitted_by="b@x")] + + with pytest.raises(ValueError, match="2 jobs are named 'j'.*a@x.*b@x.*user_name"): + find_job(jobs, "j", None, None) + + +def test_find_job_raises_when_one_name_has_two_datasites(): + jobs = [_job(datasite="do1@x"), _job(datasite="do2@x")] + + with pytest.raises(ValueError, match="2 jobs are named 'j'.*do1@x.*do2@x"): + find_job(jobs, "j", None, None) + + +def test_find_job_takes_newest_run_of_same_job(): + # An earlier run with the same submitter and datasite is still on disk. + old = _job(status="done", submitted_at="2026-09-30T09:00:00+00:00") + new = _job(status="pending", submitted_at="2026-09-30T11:00:00+00:00") + + assert find_job([old, new], "j", None, None) is new + assert find_job([new, old], "j", None, {"pending"}) is new + assert describe_job([old, new], "j", None) == "job 'j' is 'pending'" + + +def test_find_job_user_name_matches_submitter_or_datasite(): + by_a = _job(submitted_by="a@x", datasite="do@x") + at_enclave = _job(submitted_by="b@x", datasite="enclave@x") + + assert find_job([by_a, at_enclave], "j", "a@x", None) is by_a + assert find_job([by_a, at_enclave], "j", "enclave@x", None) is at_enclave + + +def test_find_job_returns_none_while_status_is_not_final(): + assert find_job([_job(status="running")], "j", None, {"done"}) is None + + +def test_find_job_accepts_any_of_several_statuses(): + job = _job(status="approved") + assert find_job([job], "j", None, {"approved", "running"}) is job + + +def test_find_job_counts_done_as_past_any_earlier_status(): + # Between two polls the job can pass the status waited for and end done. + job = _job(status="done") + + assert find_job([job], "j", None, {"approved"}) is job + assert find_job([job], "j", None, {"running", "approved"}) is job + + +def test_find_job_raises_when_job_failed_after_status(): + with pytest.raises(JobEndedError, match="ended as 'failed'"): + find_job([_job(status="failed")], "j", None, {"approved"}) + + +def test_find_job_returns_none_while_where_check_is_false(): + job = _job(status="done") + + assert find_job([job], "j", None, {"done"}, where=lambda j: False) is None + assert find_job([job], "j", None, {"done"}, where=lambda j: True) is job + + +def test_find_job_raises_on_other_final_status_with_where(): + with pytest.raises(JobEndedError, match="ended as 'rejected'"): + find_job([_job(status="rejected")], "j", None, {"done"}, where=lambda j: True) + + +def test_describe_job_names_failed_where_check(): + text = describe_job([_job(status="done")], "j", None, {"done"}, lambda j: False) + assert text == "job 'j' is 'done', and the where= check is False" + + +# ------------------------------------------------------------ find_dataset --- + + +def test_find_dataset_matches_name_and_optional_datasite(): + mine = _dataset(owner="a@x") + theirs = _dataset(owner="b@x") + + assert find_dataset([mine, _dataset(name="other")], "d", None) is mine + assert find_dataset([mine, theirs], "d", "b@x") is theirs + assert find_dataset([mine], "missing", None) is None + + +def test_find_dataset_raises_when_name_has_two_owners(): + with pytest.raises(ValueError, match="2 datasets are named 'd'.*datasite"): + find_dataset([_dataset(owner="a@x"), _dataset(owner="b@x")], "d", None) + + +# ------------------------------------------------------- waits, mock Drive --- + + +def _pair(): + return SyftRDSClient.pair_with_mock_drive_service_connection( + use_in_memory_cache=False, sync_automatically=False + ) + + +def _submit_job(ds, do, tmp_path, name="wait.job"): + code = tmp_path / "main.py" + code.write_text('with open("outputs/result.json", "w") as f:\n f.write("{}")\n') + ds.submit_python_job(user=do.email, code_path=str(code), job_name=name) + + +def _run_job(do, name="wait.job"): + do.sync() + next(j for j in do.job_client.jobs if j.name == name).approve() + do.job_runner.process_approved_jobs() + do.job_runner.share_job_results(name, share_outputs=True, share_logs=False) + do.sync() + + +def test_do_waits_until_submitted_job_arrives(tmp_path, clock): + ds, do = _pair() + _submit_job(ds, do, tmp_path) + + job = do.wait_until_has_job("wait.job", status="pending") + assert job.submitted_by == ds.email + assert clock["sleeps"] == [] + + +def test_ds_waits_until_dataset_is_shared(clock): + ds, do = _pair() + mock_path, private_path, readme_path = create_tmp_dataset_files() + + def share(): + do.create_dataset( + name="waited", + mock_path=mock_path, + private_path=private_path, + readme_path=readme_path, + summary="s", + users=[ds.email], + ) + + clock["on_sleep"].append(share) + + dataset = ds.wait_until_has_dataset("waited", datasite=do.email) + assert dataset.owner == do.email + assert clock["sleeps"] == [15] + + +def test_ds_waits_until_job_is_done_with_outputs(tmp_path, clock): + ds, do = _pair() + _submit_job(ds, do, tmp_path) + clock["on_sleep"].append(lambda: _run_job(do)) + + job = ds.wait_until_has_job( + "wait.job", status="done", where=lambda j: bool(j.output_paths) + ) + assert job.output_paths + + +def test_wait_until_has_job_timeout_names_state(clock): + ds, _ = _pair() + + with pytest.raises(TimeoutError, match="no job named 'missing'"): + ds.wait_until_has_job("missing", timeout=30) + assert clock["sleeps"] == [15, 15] + + +def test_wait_until_peered_passes_for_live_peer(clock): + ds, do = _pair() + + assert ds.wait_until_peered(do.email).email == do.email + assert clock["sleeps"] == [] diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index d5dacc3863a..8daacb29c31 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -24,7 +24,7 @@ ) from syft.sync.file_writer import FileWriter from syft.sync.job_file_change_handler import JobFileChangeHandler -from syft.sync.peers.peer import Peer +from syft.sync.peers.peer import Peer, PeerNotReadyError from syft.sync.peers.peer_list import PeerList from syft.sync.peers.peer_store import PeerStore from syft.sync.platforms.base_platform import BasePlatform @@ -45,6 +45,11 @@ DatasiteWatcherSyncerConfig, ) from syft.sync.utils.path_filters import is_normal_syncable_path +from syft.sync.utils.waiting import ( + DEFAULT_POLL_INTERVAL, + DEFAULT_WAIT_TIMEOUT, + poll_until, +) from syft.sync.utils.syftbox_utils import ( random_email, random_syftbox_folder_for_testing, @@ -479,6 +484,7 @@ class SyftboxManager(BaseModelCallbackMixin): "approve_peer_request", "reject_peer_request", "validate_peer", + "wait_until_peered", "sync", "create_checkpoint", "should_create_checkpoint", @@ -1052,6 +1058,41 @@ def validate_peer(self, peer_email: str) -> Peer: self.load_peers() return self.peer_manager.validate_peer(peer_email) + def wait_until_peered( + self, + peer_email: str, + timeout: float = DEFAULT_WAIT_TIMEOUT, + poll_interval: float = DEFAULT_POLL_INTERVAL, + ) -> Peer: + """Wait until ``validate_peer`` passes for ``peer_email``, and return the peer. + + Each poll, every ``poll_interval`` seconds, is one Drive request; the + peers load again only when that request shows a change. + + Raises: + TimeoutError: the peer did not approve within ``timeout`` seconds. + PeerSetupError: waiting cannot help, e.g. the email is not a peer. + """ + deadline = time.monotonic() + timeout + while True: + try: + return self.validate_peer(peer_email) + except PeerNotReadyError as e: + error = e + if deadline - time.monotonic() <= 0: + break + print(f"Waiting for {peer_email}: {error}") + approved = poll_until( + lambda: self.peer_manager.peer_may_be_valid(peer_email) or None, + deadline, + poll_interval, + ) + if approved is None: + break + raise TimeoutError( + f"Timed out after {timeout:g}s waiting for peer {peer_email}: {error}" + ) from error + # ========== Encryption key fingerprints ========== @property diff --git a/syft/sync/utils/waiting.py b/syft/sync/utils/waiting.py new file mode 100644 index 00000000000..dbd8245b6e5 --- /dev/null +++ b/syft/sync/utils/waiting.py @@ -0,0 +1,56 @@ +"""Poll loops for calls that wait for another party, such as a peer approval.""" + +import time +from typing import Callable, Optional, TypeVar + +T = TypeVar("T") + +# Seconds between polls. Each poll makes Drive requests; a person approving or a +# job running takes longer than this anyway. +DEFAULT_POLL_INTERVAL = 15 + +# Seconds a wait_until_* helper waits before it raises TimeoutError. +DEFAULT_WAIT_TIMEOUT = 300 + + +def poll_until( + check: Callable[[], Optional[T]], deadline: float, poll_interval: float +) -> Optional[T]: + """Sleep, then call ``check``, until it returns a value that is not None. + + ``deadline`` is a ``time.monotonic()`` value. The last sleep is cut to end + at it. Returns None when the deadline passes first. + """ + while (remaining := deadline - time.monotonic()) > 0: + time.sleep(min(poll_interval, remaining)) + if (result := check()) is not None: + return result + return None + + +def wait_for( + find: Callable[[], Optional[T]], + what: str, + describe: Callable[[], str], + timeout: float, + poll_interval: float, +) -> T: + """Return the first result of ``find`` that is not None. + + ``find`` runs at once, then after each poll. An error from ``find`` stops + the wait. ``what`` names the thing waited for; ``describe`` tells the state + now, for the messages. + + Raises: + TimeoutError: ``find`` still returns None after ``timeout`` seconds. + """ + deadline = time.monotonic() + timeout + if (found := find()) is not None: + return found + if deadline > time.monotonic(): + print(f"Waiting up to {timeout:g}s for {what}: {describe()}") + if (found := poll_until(find, deadline, poll_interval)) is not None: + return found + raise TimeoutError( + f"Timed out after {timeout:g}s waiting for {what}: {describe()}." + ) diff --git a/syft/sync/version/peer_manager.py b/syft/sync/version/peer_manager.py index 9df67c11fa3..beb478857bc 100644 --- a/syft/sync/version/peer_manager.py +++ b/syft/sync/version/peer_manager.py @@ -943,6 +943,18 @@ def forget_states_at_start(self) -> None: """ self._states_at_start = {} + def peer_may_be_valid(self, peer_email: str) -> bool: + """Tell if ``validate_peer`` can pass after the next ``load_peers()``. + + Use it between polls, after ``validate_peer`` raised + ``PeerNotReadyError``. It makes one Drive request and does not load + the peers. It is True when the peer created its folders for this + datasite, which it does when it approves our request. + """ + return any( + p.email == peer_email for p in self.connection_router.get_peer_requests() + ) + def check_peer_request_exists(self, email: str) -> bool: """Check if a peer request exists for the given email.""" return any(p.email == email for p in self.requested_by_peer_peers) diff --git a/tests/unit/test_wait_until_peered.py b/tests/unit/test_wait_until_peered.py new file mode 100644 index 00000000000..45a9049db89 --- /dev/null +++ b/tests/unit/test_wait_until_peered.py @@ -0,0 +1,105 @@ +import pytest + +import syft.sync.syftbox_manager as syftbox_manager_module +from syft.sync.connections.drive import mock_drive_service +from syft.sync.peers.peer import PeerSetupError +from syft.sync.syftbox_manager import SyftboxManager +from tests.unit.test_validate_peer import _pair, _start_both_then_request + + +def _no_sleep(_seconds): + raise AssertionError("wait_until_peered must not wait here") + + +def test_wait_until_peered_waits_for_late_approval(monkeypatch): + ds_manager, do_manager = _pair(add_peers=False) + ds_manager.add_peer(do_manager.email) + do_manager.load_peers() + waits = [] + + def approve_during_wait(seconds): + waits.append(seconds) + do_manager.approve_peer_request(ds_manager.email) + + monkeypatch.setattr(syftbox_manager_module.time, "sleep", approve_during_wait) + + peer = ds_manager.wait_until_peered(do_manager.email, timeout=60, poll_interval=1) + assert peer.is_approved + assert waits == [1] + + +def test_wait_until_peered_fails_at_once_when_waiting_cannot_help(monkeypatch): + ds_manager, do_manager = _pair(add_peers=False) + _start_both_then_request(ds_manager, do_manager) + monkeypatch.setattr(syftbox_manager_module.time, "sleep", _no_sleep) + + with pytest.raises(PeerSetupError, match="not a peer"): + ds_manager.wait_until_peered("wrong@test.com", timeout=60) + # A request that only this client can approve. + with pytest.raises(PeerSetupError, match="approve_peer_request"): + do_manager.wait_until_peered(ds_manager.email, timeout=60) + + +def _count_drive_requests(monkeypatch) -> list[int]: + """Count mock Drive requests; the returned one-item list holds the count.""" + count = [0] + for request_class in ( + mock_drive_service.MockListRequest, + mock_drive_service.MockGetRequest, + ): + original = request_class.execute + + def counted(self, *args, _original=original, **kwargs): + count[0] += 1 + return _original(self, *args, **kwargs) + + monkeypatch.setattr(request_class, "execute", counted) + # Downloads go through get_media(), not execute(). + original_get_media = mock_drive_service.MockFilesResource.get_media + + def counted_get_media(self, *args, **kwargs): + count[0] += 1 + return original_get_media(self, *args, **kwargs) + + monkeypatch.setattr( + mock_drive_service.MockFilesResource, "get_media", counted_get_media + ) + return count + + +def _count_load_peers(monkeypatch) -> list[int]: + count = [0] + original = SyftboxManager.load_peers + + def counted(self, *args, **kwargs): + count[0] += 1 + return original(self, *args, **kwargs) + + monkeypatch.setattr(SyftboxManager, "load_peers", counted) + return count + + +def test_wait_until_peered_polls_with_one_drive_request(monkeypatch): + ds_manager, do_manager = _pair(add_peers=False) + ds_manager.add_peer(do_manager.email) + drive_requests = _count_drive_requests(monkeypatch) + loads = _count_load_peers(monkeypatch) + clock = [0.0] + per_poll = [] + + def record_poll(seconds): + clock[0] += seconds + per_poll.append(drive_requests[0]) + drive_requests[0] = 0 + + monkeypatch.setattr(syftbox_manager_module.time, "sleep", record_poll) + monkeypatch.setattr(syftbox_manager_module.time, "monotonic", lambda: clock[0]) + + with pytest.raises(TimeoutError, match="not approved"): + ds_manager.wait_until_peered(do_manager.email, timeout=4, poll_interval=1) + # Four polls in four seconds. Before the first sleep: the initial full + # load. After it: one Drive request per poll, and no other load. + assert len(per_poll) == 4 + assert per_poll[1:] == [1, 1, 1] + assert drive_requests[0] == 1 + assert loads[0] == 1 diff --git a/tests/unit/test_waiting.py b/tests/unit/test_waiting.py new file mode 100644 index 00000000000..efcd9932221 --- /dev/null +++ b/tests/unit/test_waiting.py @@ -0,0 +1,76 @@ +import time + +import pytest + +from syft.sync.utils.waiting import poll_until, wait_for + + +@pytest.fixture +def clock(monkeypatch): + """Fake monotonic clock; time.sleep() advances it and records each sleep.""" + state = {"now": 0.0, "sleeps": []} + + def sleep(seconds): + state["sleeps"].append(seconds) + state["now"] += seconds + + monkeypatch.setattr(time, "sleep", sleep) + monkeypatch.setattr(time, "monotonic", lambda: state["now"]) + return state + + +def test_poll_until_returns_first_result_not_none(clock): + results = iter([None, None, "found"]) + + assert poll_until(lambda: next(results), deadline=100, poll_interval=10) == "found" + assert clock["sleeps"] == [10, 10, 10] + + +def test_poll_until_cuts_last_sleep_at_deadline(clock): + assert poll_until(lambda: None, deadline=25, poll_interval=10) is None + assert clock["sleeps"] == [10, 10, 5] + + +def test_wait_for_returns_at_once_when_found(clock): + assert wait_for(lambda: 1, "x", lambda: "", timeout=60, poll_interval=10) == 1 + assert clock["sleeps"] == [] + + +def test_wait_for_polls_until_found(clock): + results = iter([None, None, "found"]) + + found = wait_for( + lambda: next(results), "x", lambda: "", timeout=60, poll_interval=10 + ) + assert found == "found" + assert clock["sleeps"] == [10, 10] + + +def test_wait_for_timeout_names_target_and_state(clock): + with pytest.raises(TimeoutError, match="job 'j'.*job 'j' is 'pending'"): + wait_for( + lambda: None, + "job 'j'", + lambda: "job 'j' is 'pending'", + timeout=30, + poll_interval=10, + ) + assert clock["sleeps"] == [10, 10, 10] + + +def test_wait_for_with_zero_timeout_checks_once(clock): + calls = [] + + with pytest.raises(TimeoutError): + wait_for(lambda: calls.append(1), "x", lambda: "", timeout=0, poll_interval=10) + assert calls == [1] + assert clock["sleeps"] == [] + + +def test_wait_for_raises_find_error_at_once(clock): + def find(): + raise ValueError("ambiguous") + + with pytest.raises(ValueError, match="ambiguous"): + wait_for(find, "x", lambda: "", timeout=60, poll_interval=10) + assert clock["sleeps"] == []