Auto-track and globally shut down all Forge actors and services #357

DNXie · 2025-10-09T00:00:49Z

Context: #360

This PR introduces centralized allocation tracking and unified shutdown for all ForgeActor and ServiceInterface instances managed by the global Provisioner.

Actors and services are now automatically registered with the provisioner when spawned via .as_actor() or .as_service().
This enables a single, global shutdown sequence (await shutdown()) that gracefully terminates all registered allocations in reverse order, no manual teardown required.

Automatic Registration

Added register_service() / register_actor() to the Provisioner with type checks.
Added matching top-level helpers for API consistency (get_proc_mesh, stop_proc_mesh).
ForgeActor.as_service() and .as_actor() now automatically register their proxies after initialization.

Centralized Shutdown

Added Provisioner.shutdown_all_allocations() for unified teardown.
Gracefully shuts down all tracked ServiceInterface and ForgeActor instances in reverse allocation order.
Move metric logger shutdown into the centralized function shutdown()

Test:

python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml

The log after ctrl+C (cleaned version):

Shutting down...
Shutting down metric logger...
Shutting down provisioner..
Shutting down 3 service(s) and 4 actor(s)...
Health loop stopped gracefully.
Health loop stopped gracefully.
Shutdown completed successfully

DNXie · 2025-10-09T19:07:05Z

src/forge/controller/actor.py

        This method is used by `Service` to teardown a replica.
        """
+        if not quiet:
+            logger.info(f"Shutting down actor {getattr(actor, 'name', cls.__name__)}")


Adding this quiet check because otherwise when we shutdown a service, it would call actor.shutdown and print the log twice.

allenwang28 · 2025-10-10T15:17:48Z

src/forge/controller/provisioner.py


+    async def track_allocation(self, alloc: Any):
+        """Tracks an allocation for cleanup."""
+        self._allocations.append(alloc)


hmmm, I think an even simpler approach is to just track the proc meshes right? We can just do await proc_mesh.stop() and I think everything inside of it should shut down neatly. Let me know if that doesn't work though

I think it would work for actor since shutting down actor is essentially stopping the proc_mesh.
But for service, it involves some other operations such as stopping the replicas and healthy loop.

allenwang28

ok, I understand - here are a few suggestions:

In provisioner.py, let's add in a register_service and register_actor on the Provisioner
Also add top level register_service and register_actor functions (like get_proc_mesh)
Within .as_service(...) and .as_actor(...) we call register_service and register_actor respectively

DNXie · 2025-10-13T18:10:38Z

ok, I understand - here are a few suggestions:

In provisioner.py, let's add in a register_service and register_actor on the Provisioner

Also add top level register_service and register_actor functions (like get_proc_mesh)

Within .as_service(...) and .as_actor(...) we call register_service and register_actor respectively

@allenwang28 Done!

src/forge/controller/actor.py

src/forge/controller/provisioner.py

allenwang28 · 2025-10-13T18:28:42Z

src/forge/controller/provisioner.py

+
+    def register_actor(self, actor: "ForgeActor") -> None:
+        """Registers a single actor allocation for cleanup."""
+        from monarch._src.actor.actor_mesh import ActorMesh


can we do this at the toplevel? And use

from monarch.actor import ActorMesh

I got

ImportError: cannot import name 'ActorMesh' from 'monarch.actor' (/home/dxie/.fbpkg_conda_envs/forge-e146614/lib/python3.10/site-packages/monarch/actor/__init__.py)

Sure I can do it on top-level.

allenwang28 · 2025-10-13T18:29:26Z

src/forge/controller/service/interface.py

        """
        Shut down the underlying Service.
        """
+        logger.info(f"Shutting down service {self.actor_def.__name__}")


since we're quietly shutting down actors/replicas, in this log here can we mention how many actors we're shutting down?

Done. Now the log only says

Shutting down 3 service(s) and 4 actor(s)...

…ovisioner_shutdown

allenwang28 · 2025-10-13T19:50:19Z

apps/grpo/main.py

-
        # give mlogger time to shutdown backends, otherwise they can stay running.
        # TODO (felipemello) find more elegant solution
        await mlogger.shutdown.call_one()


@DNXie @felipemello1 maybe we can just move the mlogger shutdown into the global shutdown as well?

I moved it into shutdown().

allenwang28

thanks Danning!

allenwang28 · 2025-10-14T21:46:54Z

src/forge/controller/provisioner.py


+async def shutdown_metric_logger():
+    """Shutdown the global metric logger and all its backends."""
+    from forge.observability.metric_actors import get_or_create_metric_logger


if this doesn't need to be inline imported, can we do this at the toplevel? This should be a general thing

This will cause a circular dependency issue. This is also what Felipe did: https://github.com/meta-pytorch/forge/blob/main/src/forge/controller/provisioner.py#L325-L327

ok sg! I will just assume this to be the case moving forward

codecov-commenter · 2025-10-14T21:54:02Z

Codecov Report

❌ Patch coverage is 46.15385% with 28 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@4c14792). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
src/forge/controller/provisioner.py	39.13%	28 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #357   +/-   ##
=======================================
  Coverage        ?   64.63%           
=======================================
  Files           ?       79           
  Lines           ?     7736           
  Branches        ?        0           
=======================================
  Hits            ?     5000           
  Misses          ?     2736           
  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.

…-pytorch#357)

centralized shutdown

cc5fb5a

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

DNXie changed the title ~~Auto-track and globally shut down all Forge actors and services~~ [WIP] Auto-track and globally shut down all Forge actors and services Oct 9, 2025

DNXie added 4 commits October 8, 2025 18:16

debug. TODO: simplify

9cc5531

debug

84dcd17

update all main

8ba3e5f

simplify shutting down

5b8d81c

DNXie commented Oct 9, 2025

View reviewed changes

clean up code

5fe99d5

DNXie changed the title ~~[WIP] Auto-track and globally shut down all Forge actors and services~~ Auto-track and globally shut down all Forge actors and services Oct 9, 2025

DNXie requested review from allenwang28 and felipemello1 October 9, 2025 19:18

DNXie marked this pull request as ready for review October 9, 2025 19:18

Merge branch 'main' into provisioner_shutdown

36225af

allenwang28 reviewed Oct 10, 2025

View reviewed changes

fix lint and update docstring

d705df9

DNXie requested a review from allenwang28 October 10, 2025 23:45

allenwang28 reviewed Oct 13, 2025

View reviewed changes

DNXie added 3 commits October 13, 2025 11:06

add function register_xx

52fa867

rollback changes

f0ba99a

cleanup

4e477b5

DNXie requested a review from allenwang28 October 13, 2025 18:10

allenwang28 reviewed Oct 13, 2025

View reviewed changes

DNXie and others added 2 commits October 13, 2025 11:40

resolve comment

e7ae73d

Merge branch 'main' into provisioner_shutdown

eb9f667

DNXie requested a review from allenwang28 October 13, 2025 18:49

DNXie added 3 commits October 13, 2025 12:05

fix bug

77c2344

Merge branch 'provisioner_shutdown' of github.com:DNXie/forge into pr…

b39fc4b

…ovisioner_shutdown

fix lint

d4f3d57

allenwang28 reviewed Oct 13, 2025

View reviewed changes

DNXie and others added 3 commits October 13, 2025 16:59

Merge branch 'main' into provisioner_shutdown

0db9ed6

move mlogger.shutdown into global shutdown()

f31aa3a

cleanup, update other main

ffbf0ca

DNXie requested a review from allenwang28 October 13, 2025 23:03

allenwang28 approved these changes Oct 14, 2025

View reviewed changes

Merge branch 'main' into provisioner_shutdown

7c5417d

fix broken ci

5dc138c

DNXie merged commit f19eb1b into meta-pytorch:main Oct 14, 2025
9 checks passed

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

Auto-track and globally shut down all Forge actors and services (meta…

8ae6743

…-pytorch#357)

DNXie mentioned this pull request Oct 16, 2025

Problems with Clean Shutdown and Actor Management in training scripts #360

Open

Auto-track and globally shut down all Forge actors and services #357

Auto-track and globally shut down all Forge actors and services #357

Uh oh!

Conversation

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

Uh oh!

Automatic Registration

Centralized Shutdown

Test:

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

allenwang28 left a comment

Choose a reason for hiding this comment

Uh oh!

DNXie commented Oct 13, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

allenwang28 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Oct 14, 2025

Codecov Report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

DNXie commented Oct 9, 2025 •

edited

Loading

DNXie Oct 13, 2025 •

edited

Loading

DNXie Oct 13, 2025 •

edited

Loading