feat(fragile): enable multiprocessing solver by default

 - fragile == 0.0.47
 - this is almost always going to be faster because mathy environments aren't optimized for speed to begin with, so they max out a single-process python app quickly when it comes to env iterations

Author: justindujardin

libraries/mathy_python/mathy/agents/fragile.py

@@ -13,6 +13,7 @@
 from fragile.core.swarm import Swarm
 from fragile.core.utils import StateDict
 from fragile.core.walkers import Walkers
+from fragile.distributed.env import ParallelEnv
 from gym import spaces
 from pydantic import BaseModel
 from wasabi import msg
@@ -29,6 +30,7 @@
 
 
 class SwarmConfig(BaseModel):
+    use_mp: bool = True
     verbose: bool = False
     n_walkers: int = 512
     max_iters: int = 100
@@ -193,6 +195,8 @@ def swarm_solve(problem: str, config: SwarmConfig):
         name="mathy_v0", problem=problem, repeat_problem=True
     )
     mathy_env: MathyEnv = env_callable()._env._env.mathy
+    if config.use_mp:
+        env_callable = ParallelEnv(env_callable=env_callable)
     swarm = Swarm(
         model=lambda env: DiscreteMasked(env=env),
         env=env_callable,

libraries/mathy_python/mathy/cli.py

@@ -36,6 +36,13 @@ def cli_contribute():
     is_flag=True,
     help="Use swarm solver from fragile library without a trained model",
 )
+@click.option(
+    "parallel",
+    "--parallel",
+    default=True,
+    is_flag=True,
+    help="Use parallel execution with the swarm solver",
+)
 @click.option(
     "model", "--model", default="mathy_alpha_sm", help="The path to a mathy model",
 )
@@ -47,7 +54,9 @@ def cli_contribute():
     help="The max number of steps before the episode is over",
 )
 @click.argument("problem", type=str)
-def cli_simplify(agent: str, problem: str, model: str, max_steps: int, swarm: bool):
+def cli_simplify(
+    agent: str, problem: str, model: str, max_steps: int, swarm: bool, parallel: bool
+):
     """Simplify an input polynomial expression."""
     setup_tf_env()
 
@@ -57,12 +66,12 @@ def cli_simplify(agent: str, problem: str, model: str, max_steps: int, swarm: bo
 
     mt: Mathy
     if swarm is True:
-        mt = Mathy(config=SwarmConfig())
+        mt = Mathy(config=SwarmConfig(use_mp=parallel))
     else:
         try:
             mt = load_model(model)
         except ValueError:
-            mt = Mathy(config=SwarmConfig())
+            mt = Mathy(config=SwarmConfig(use_mp=parallel))
 
     mt.simplify(problem=problem, max_steps=max_steps)
 

libraries/mathy_python/setup.py

@@ -23,7 +23,7 @@ def setup_package():
         DEVELOPMENT_MODULES = [line.strip() for line in file if "-e" not in line]
 
     extras = {
-        "fragile": ["fragile==0.0.45"],
+        "fragile": ["fragile==0.0.47"],
         "dev": DEVELOPMENT_MODULES,
     }
     extras["all"] = [item for group in extras.values() for item in group]

libraries/mathy_python/tests/test_api.py

@@ -24,11 +24,6 @@ def test_mathy_with_model_and_config():
     mt = Mathy(model=model, config=config)
 
 
-def test_api_mathy_simplify():
-    mt: Mathy = Mathy(config=SwarmConfig())
-    mt.simplify(problem="2x+4x", max_steps=4)
-
-
 def test_api_mathy_constructor():
     # Defaults to swarm config
     assert isinstance(Mathy().state, MathyAPISwarmState)

libraries/mathy_python/tests/test_cli.py

@@ -32,9 +32,13 @@ def test_cli_simplify():
         assert result.exit_code == 0
 
 
-def test_cli_simplify_swarm():
+@pytest.mark.parametrize("use_mp", [True, False])
+def test_cli_simplify_swarm(use_mp: bool):
     runner = CliRunner()
-    result = runner.invoke(cli, ["simplify", "4x + 2x", "--swarm"])
+    args = ["simplify", "4x + 2x", "--swarm"]
+    if use_mp:
+        args.append("--parallel")
+    result = runner.invoke(cli, args)
     assert result.exit_code == 0