Fix policy update test #365

JenniferWang · 2025-10-09T21:19:35Z

Fix the test fixture yaml files
Re-write the logic to make the round-trip test compatible with multi-nodes.

Specifically,

Move validation function to the policy.py module because @endpoint almost only work with lambda functions.
Refactor the test to make setup / tear down clearer

Closes sub-task in #411, close #143

Jack-Khuu · 2025-10-09T21:34:01Z

Can we close #143 once this lands too :)?

casteryh · 2025-10-09T22:37:20Z

tests/integration_tests/test_policy_update.py

+    )
+
+    # Cleanup DCP directory
+    path = Path(TEST_DCP_DIR)


You can use TemporaryDirectory which handles this automatically

I have to change this because NFS does not work for tmp directory

casteryh · 2025-10-09T22:43:16Z

src/forge/actors/policy.py

    # We only care about the final output
    params.output_kind = RequestOutputKind.FINAL_ONLY
    return params
+


Can you explain why this has to be in policy.py? Is monarch failing to pickup/serialize the function if it's not in here?
If it has to be here, may be prepend all the functions with _.

Yes, monarch pickled the function with the module path integration_tests but this cannot be resolved in the remote node. Seems that we need to define it in the main modules.

codecov-commenter · 2025-10-10T01:41:29Z

Codecov Report

❌ Patch coverage is 0% with 126 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@2a6e46f). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
tests/integration_tests/test_policy_update.py	0.00%	82 Missing ⚠️
src/forge/actors/policy.py	0.00%	44 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #365   +/-   ##
=======================================
  Coverage        ?   63.24%           
=======================================
  Files           ?       78           
  Lines           ?     7725           
  Branches        ?        0           
=======================================
  Hits            ?     4886           
  Misses          ?     2839           
  Partials        ?        0

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

joecummings

While the test is amazing, having it in the policy.py is not really an option IMO.

From a UX / library perspective, it makes things way more difficult to understand. The argument could be made that directing users to docs alleviates this pain, but 1) our docs are incredibly incomplete and 2) we know from OSS users experience with titan/ tune that most people prefer to just go to the code.

allenwang28 · 2025-10-10T15:12:56Z

While the test is amazing, having it in the policy.py is not really an option IMO.

Oh I completely missed that part. Is it possible to create an inherited Policy that has all of the test functionality and is kept in integration_tests?

I'm not sure I'm clear what the monarch pickling issue is

joecummings · 2025-10-10T15:13:24Z

@allenwang28 What are our options here? Surely someone else has tried to run a multi node test via Monarch and not have to clutter up an Actor?

Jack-Khuu · 2025-10-10T16:10:06Z

Wild thought (lmk if this errors): For tests specific methods, can we just append them to the class definition in the test? It's the same idea as inheriting, but let's you surgically change the nested PolicyWorker

Within test:

from foo import Bar

def util(self):
  ...

Bar._special_test_util = util

proxy: Bar = Bar()

I'm not sure I'm clear what the monarch pickling issue is

I'd like to see the trace too since at least 2 others have encountered this when using Monarch/Services

JenniferWang · 2025-10-13T19:28:39Z

@joecummings , @allenwang28 , @Jack-Khuu

The pickle problem is, what I think, caused by how Monarch serializes arguments (in this case, passing a lambda v.s. a function variable to an endpoint on the Policy service.

WARNING  forge.controller.service.replica:replica.py:257 Got failure on replica 0. Error:
A remote actor call has failed.
 Traceback of where the remote call failed (most recent call last):
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/actor_mesh.py", line 875, in handle
    result = await instrumented()
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/actor_mesh.py", line 872, in instrumented
    raise e
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/actor_mesh.py", line 865, in instrumented
    result = await the_method(*args, **kwargs)
  File "/mnt/home/jiyue/forge/src/forge/actors/policy.py", line 451, in _test_validate_model_params
    return await self.policy_worker._test_validate_model_params.call(validate_fn)
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/future.py", line 138, in mark_complete
    func, value = fut.set_result, await coro
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/endpoint.py", line 147, in process
    rank, value = await r._recv()
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/actor_mesh.py", line 728, in _recv
    return self._process(result)
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/actor_mesh.py", line 758, in _process
    return rank, super()._process(msg)
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/actor_mesh.py", line 737, in _process
    raise cast(Exception, payload)
monarch._src.actor.actor_mesh.ActorError: A remote actor call has failed.
 Traceback of where the remote call failed (most recent call last):
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/actor_mesh.py", line 811, in handle
    args, kwargs = unflatten(message, local_state)
  File "/mnt/data/jiyue/lib/python3.10/site-packages/monarch/_src/actor/pickle.py", line 98, in unflatten
    return up.load()
ModuleNotFoundError: No module named 'tests.integration_tests'

allenwang28 · 2025-10-13T19:36:17Z

@JenniferWang how do you run this workload?

JenniferWang · 2025-10-13T20:35:34Z

@JenniferWang how do you run this workload?

pytest -s tests/integration_tests/test_policy_update.py::TestWeightSync::test_sanity_check         --config apps/grpo/qwen3_1_7b.yaml

JenniferWang · 2025-10-13T21:08:40Z

@allenwang28 , indeed, i think it's because of the way how the script is invoked.

allenwang28 · 2025-10-13T21:19:08Z

got it, I wonder if PYTHONPATH=. pytest ... works?

JenniferWang · 2025-10-14T13:24:26Z

got it, I wonder if PYTHONPATH=. pytest ... works?

Magic!

casteryh · 2025-10-14T17:19:29Z

got it, I wonder if PYTHONPATH=. pytest ... works?

Magic!

Can you add how to run this test to the docstring?

tests/integration_tests/test_policy_update.py

meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 9, 2025

JenniferWang requested review from allenwang28 and casteryh October 9, 2025 21:22

JenniferWang marked this pull request as ready for review October 9, 2025 21:23

casteryh reviewed Oct 9, 2025

View reviewed changes

casteryh approved these changes Oct 10, 2025

View reviewed changes

allenwang28 approved these changes Oct 10, 2025

View reviewed changes

joecummings requested changes Oct 10, 2025

View reviewed changes

JenniferWang force-pushed the fix-policy-update-test branch 2 times, most recently from 290e506 to 13c8499 Compare October 14, 2025 13:23

JenniferWang requested review from allenwang28 and joecummings October 14, 2025 13:24

JenniferWang commented Oct 14, 2025

View reviewed changes

tests/integration_tests/test_policy_update.py Show resolved Hide resolved

JenniferWang added 2 commits October 15, 2025 09:23

rebase + squash

917051c

new api

bf0289e

JenniferWang force-pushed the fix-policy-update-test branch from fb47d9b to bf0289e Compare October 15, 2025 13:48

joecummings approved these changes Oct 15, 2025

View reviewed changes

JenniferWang merged commit 703f419 into main Oct 15, 2025
9 checks passed

allenwang28 pushed a commit to allenwang28/forge that referenced this pull request Oct 15, 2025

Fix policy update test (meta-pytorch#365)

f580316

Fix policy update test #365

Fix policy update test #365

Uh oh!

Conversation

JenniferWang commented Oct 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Jack-Khuu commented Oct 9, 2025

Uh oh!

casteryh Oct 9, 2025

Choose a reason for hiding this comment

Uh oh!

JenniferWang Oct 10, 2025

Choose a reason for hiding this comment

Uh oh!

casteryh Oct 9, 2025

Choose a reason for hiding this comment

Uh oh!

JenniferWang Oct 10, 2025

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Oct 10, 2025

Codecov Report

Uh oh!

joecummings left a comment

Choose a reason for hiding this comment

Uh oh!

allenwang28 commented Oct 10, 2025

Uh oh!

joecummings commented Oct 10, 2025

Uh oh!

Jack-Khuu commented Oct 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JenniferWang commented Oct 13, 2025

Uh oh!

allenwang28 commented Oct 13, 2025

Uh oh!

JenniferWang commented Oct 13, 2025

Uh oh!

JenniferWang commented Oct 13, 2025

Uh oh!

allenwang28 commented Oct 13, 2025

Uh oh!

JenniferWang commented Oct 14, 2025

Uh oh!

casteryh commented Oct 14, 2025

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

JenniferWang commented Oct 9, 2025 •

edited

Loading

Jack-Khuu commented Oct 10, 2025 •

edited

Loading