[monarch] envvar to force v1 by default

Having monarch.actor selectively import from the same parts. This fixes the public API. Internal uses will probably need to call the add create_host_mesh_from_alloc which will give the right kind of host mesh based on environment variable.

`MONARCH_HOST_MESH_V1_REMOVE_ME_BEFORE_RELEASE` is the name of the flag. Chosen because this needs to be gone before the public release.

Differential Revision: [D84107651](https://our.internmc.facebook.com/intern/diff/D84107651/)

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D84107651/)!

ghstack-source-id: 314700169
Pull Request resolved: #1463

Author: zdevito

python/monarch/_src/actor/v1/__init__.py

@@ -5,3 +5,28 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-unsafe
+import os
+
+from monarch._rust_bindings.monarch_hyperactor.alloc import AllocConstraints
+from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice
+
+from monarch._src.actor.allocator import AllocateMixin
+
+from monarch._src.actor.endpoint import Extent
+from monarch._src.actor.host_mesh import HostMesh as HostMeshV0
+from monarch._src.actor.v1.host_mesh import HostMesh as HostMeshV1
+
+enabled = os.environ.get("MONARCH_HOST_MESH_V1_REMOVE_ME_BEFORE_RELEASE", "0") != "0"
+
+
+def host_mesh_from_alloc(
+    name: str, extent: Extent, allocator: AllocateMixin, constraints: AllocConstraints
+) -> "HostMeshV0 | HostMeshV1":
+    if enabled:
+        return HostMeshV1.allocate_nonblocking(name, extent, allocator, constraints)
+    else:
+        return HostMeshV0(
+            Shape(extent.labels, Slice.new_row_major(extent.sizes)),
+            allocator,
+            constraints,
+        )

python/monarch/_src/job/job.py

@@ -14,8 +14,7 @@
 
 # note: the jobs api is intended as a library so it should
 # only be importing _public_ monarch API functions.
-from monarch._src.actor.host_mesh import HostMesh, this_host
-from typing_extensions import Self
+from monarch.actor import HostMesh, this_host
 
 
 class JobState:

python/monarch/_src/job/meta.py

@@ -19,14 +19,14 @@
 )
 
 from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask, Shared
-from monarch._rust_bindings.monarch_hyperactor.shape import Shape, Slice
+from monarch._rust_bindings.monarch_hyperactor.shape import Extent
 from monarch._src.actor.allocator import AllocateMixin
-from monarch._src.actor.host_mesh import HostMesh
 from monarch._src.actor.meta.allocator import (
     MastAllocator,
     MastAllocatorBase,
     MastAllocatorConfig,
 )
+from monarch._src.actor.v1 import host_mesh_from_alloc
 
 from monarch._src.job.job import BatchJob, JobState, JobTrait
 
@@ -173,10 +173,8 @@ def _state(self) -> JobState:
                 job_started,
             )
             constraints = AllocConstraints({MastAllocator.ALLOC_LABEL_TASK_GROUP: name})
-            host_meshes[name] = HostMesh(
-                Shape(["hosts"], Slice.new_row_major([num_host])),
-                allocator,
-                constraints,
+            host_meshes[name] = host_mesh_from_alloc(
+                name, Extent(["hosts"], [num_host]), allocator, constraints
             )
 
         return JobState(host_meshes)

python/monarch/actor/__init__.py

@@ -9,6 +9,8 @@
 Monarch Actor API - Public interface for actor functionality.
 """
 
+from typing import TYPE_CHECKING
+
 from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport
 from monarch._rust_bindings.monarch_hyperactor.shape import Extent
 from monarch._src.actor.actor_mesh import (
@@ -34,19 +36,17 @@
 from monarch._src.actor.endpoint import endpoint
 from monarch._src.actor.future import Future
 
-from monarch._src.actor.host_mesh import (
-    HostMesh,
-    hosts_from_config,
-    this_host,
-    this_proc,
-)
-from monarch._src.actor.proc_mesh import (
-    get_or_spawn_controller,
-    local_proc_mesh,
-    proc_mesh,
-    ProcMesh,
-    sim_proc_mesh,
-)
+from monarch._src.actor.host_mesh import hosts_from_config
+from monarch._src.actor.proc_mesh import local_proc_mesh, proc_mesh, sim_proc_mesh
+
+from monarch._src.actor.v1 import enabled as v1_enabled
+
+if TYPE_CHECKING or not v1_enabled:
+    from monarch._src.actor.host_mesh import HostMesh, this_host, this_proc
+    from monarch._src.actor.proc_mesh import get_or_spawn_controller, ProcMesh
+else:
+    from monarch._src.actor.v1.host_mesh import HostMesh, this_host, this_proc
+    from monarch._src.actor.v1.proc_mesh import get_or_spawn_controller, ProcMesh
 
 __all__ = [
     "Accumulator",