Refactoring deployment a bit (#15)

examples/cloud_deployment/gcp/deploy.py

@@ -1,23 +1,13 @@
 import getpass
-import os
 
-from prediction_market_agent_tooling.deploy.gcp.deploy import (
-    deploy_to_gcp,
-    remove_deployed_gcp_function,
-    run_deployed_gcp_function,
-    schedule_deployed_gcp_function,
-)
-from prediction_market_agent_tooling.deploy.gcp.utils import gcp_function_is_active
+from prediction_market_agent_tooling.deploy.agent_example import DeployableCoinFlipAgent
 from prediction_market_agent_tooling.markets.markets import MarketType
 
 if __name__ == "__main__":
-    current_dir = os.path.dirname(os.path.realpath(__file__))
-    fname = deploy_to_gcp(
-        requirements_file=None,
-        extra_deps=[
-            "git+https://github.com/gnosis/prediction-market-agent-tooling.git@main"
-        ],
-        function_file=f"{current_dir}/agent.py",
+    agent = DeployableCoinFlipAgent()
+    agent.deploy_gcp(
+        # TODO: Switch to main.
+        repository="git+https://github.com/gnosis/prediction-market-agent-tooling.git@peter/refactor-deployment",
         market_type=MarketType.MANIFOLD,
         labels={
             "owner": getpass.getuser()
@@ -27,18 +17,6 @@
         secrets={
             "MANIFOLD_API_KEY": f"JUNG_PERSONAL_GMAIL_MANIFOLD_API_KEY:latest"
         },  # Must be in the format "env_var_in_container => secret_name:version", you can create secrets using `gcloud secrets create --labels owner=<your-name> <secret-name>` command.
-        memory=512,
+        memory=256,
+        cron_schedule="0 */2 * * *",
     )
-
-    # Check that the function is deployed
-    assert gcp_function_is_active(fname)
-
-    # Run the function
-    response = run_deployed_gcp_function(fname)
-    assert response.ok
-
-    # Schedule the function to run once every 2 hours
-    schedule_deployed_gcp_function(fname, cron_schedule="0 */2 * * *")
-
-    # Delete the function
-    remove_deployed_gcp_function(fname)

prediction_market_agent_tooling/deploy/agent.py

@@ -1,9 +1,16 @@
+import inspect
+import os
+import tempfile
 import time
+import typing as t
 from decimal import Decimal
-from enum import Enum
-
-from pydantic import BaseModel
 
+from prediction_market_agent_tooling.deploy.gcp.deploy import (
+    deploy_to_gcp,
+    run_deployed_gcp_function,
+    schedule_deployed_gcp_function,
+)
+from prediction_market_agent_tooling.deploy.gcp.utils import gcp_function_is_active
 from prediction_market_agent_tooling.markets.data_models import (
     AgentMarket,
     BetAmount,
@@ -16,12 +23,19 @@
 )
 
 
-class DeploymentType(str, Enum):
-    GOOGLE_CLOUD = "google_cloud"
-    LOCAL = "local"
+class DeployableAgent:
+    def __init__(self) -> None:
+        self.load()
+
+    def __init_subclass__(cls, **kwargs: t.Any) -> None:
+        if cls.__init__ is not DeployableAgent.__init__:
+            raise TypeError(
+                "Cannot override __init__ method of DeployableAgent class, please override the `load` method to set up the agent."
+            )
 
+    def load(self) -> None:
+        pass
 
-class DeployableAgent(BaseModel):
     def pick_markets(self, markets: list[AgentMarket]) -> list[AgentMarket]:
         """
         This method should be implemented by the subclass to pick the markets to bet on. By default, it picks only the first market.
@@ -34,26 +48,72 @@ def answer_binary_market(self, market: AgentMarket) -> bool:
         """
         raise NotImplementedError("This method must be implemented by the subclass")
 
-    def deploy(
+    def deploy_local(
         self,
         market_type: MarketType,
-        deployment_type: DeploymentType,
         sleep_time: float,
         timeout: float,
         place_bet: bool,
     ) -> None:
-        if deployment_type == DeploymentType.GOOGLE_CLOUD:
-            # Deploy to Google Cloud Functions, and use Google Cloud Scheduler to run the function
-            raise NotImplementedError(
-                "TODO not currently possible via DeployableAgent class. See examples/cloud_deployment/ instead."
+        start_time = time.time()
+        while True:
+            self.run(market_type=market_type, _place_bet=place_bet)
+            time.sleep(sleep_time)
+            if time.time() - start_time > timeout:
+                break
+
+    def deploy_gcp(
+        self,
+        repository: str,
+        market_type: MarketType,
+        memory: int,
+        labels: dict[str, str] | None = None,
+        env_vars: dict[str, str] | None = None,
+        secrets: dict[str, str] | None = None,
+        cron_schedule: str | None = None,
+    ) -> None:
+        path_to_agent_file = os.path.relpath(inspect.getfile(self.__class__))
+
+        entrypoint_template = f"""
+from {path_to_agent_file.replace("/", ".").replace(".py", "")} import *
+import functions_framework
+from prediction_market_agent_tooling.markets.markets import MarketType
+
+@functions_framework.http
+def main(request) -> str:
+    {self.__class__.__name__}().run(market_type={market_type.__class__.__name__}.{market_type.name})
+    return "Success"
+"""
+
+        gcp_fname = self.get_gcloud_fname(market_type)
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+            f.write(entrypoint_template)
+            f.flush()
+
+            fname = deploy_to_gcp(
+                gcp_fname=gcp_fname,
+                requirements_file=None,
+                extra_deps=[repository],
+                function_file=f.name,
+                labels=labels,
+                env_vars=env_vars,
+                secrets=secrets,
+                memory=memory,
             )
-        elif deployment_type == DeploymentType.LOCAL:
-            start_time = time.time()
-            while True:
-                self.run(market_type=market_type, _place_bet=place_bet)
-                time.sleep(sleep_time)
-                if time.time() - start_time > timeout:
-                    break
+
+        # Check that the function is deployed
+        if not gcp_function_is_active(fname):
+            raise RuntimeError("Failed to deploy the function")
+
+        # Run the function
+        response = run_deployed_gcp_function(fname)
+        if not response.ok:
+            raise RuntimeError("Failed to run the deployed function")
+
+        # Schedule the function
+        if cron_schedule:
+            schedule_deployed_gcp_function(fname, cron_schedule=cron_schedule)
 
     def run(self, market_type: MarketType, _place_bet: bool = True) -> None:
         available_markets = [
@@ -71,9 +131,8 @@ def run(self, market_type: MarketType, _place_bet: bool = True) -> None:
                     omen_auto_deposit=True,
                 )
 
-    @classmethod
-    def get_gcloud_fname(cls, market_type: MarketType) -> str:
-        return f"{cls.__class__.__name__.lower()}-{market_type}-{int(time.time())}"
+    def get_gcloud_fname(self, market_type: MarketType) -> str:
+        return f"{self.__class__.__name__.lower()}-{market_type}-{int(time.time())}"
 
 
 def get_tiny_bet(market_type: MarketType) -> BetAmount:

examples/cloud_deployment/gcp/agent.py → prediction_market_agent_tooling/deploy/agent_example.py

@@ -1,11 +1,7 @@
 import random
 
-import functions_framework
-from flask.wrappers import Request
-
 from prediction_market_agent_tooling.deploy.agent import DeployableAgent
 from prediction_market_agent_tooling.markets.data_models import AgentMarket
-from prediction_market_agent_tooling.markets.markets import MarketType
 
 
 class DeployableCoinFlipAgent(DeployableAgent):
@@ -16,9 +12,3 @@ def pick_markets(self, markets: list[AgentMarket]) -> list[AgentMarket]:
 
     def answer_binary_market(self, market: AgentMarket) -> bool:
         return random.choice([True, False])
-
-
-@functions_framework.http
-def main(request: Request) -> str:
-    DeployableCoinFlipAgent().run(market_type=MarketType.MANIFOLD)
-    return "Success"

prediction_market_agent_tooling/deploy/gcp/deploy.py

@@ -7,7 +7,6 @@
 import requests
 from cron_validator import CronValidator
 
-from prediction_market_agent_tooling.deploy.agent import DeployableAgent
 from prediction_market_agent_tooling.deploy.gcp.utils import (
     gcloud_create_topic_cmd,
     gcloud_delete_function_cmd,
@@ -17,18 +16,17 @@
     get_gcloud_function_uri,
     get_gcloud_id_token,
 )
-from prediction_market_agent_tooling.markets.markets import MarketType
 from prediction_market_agent_tooling.tools.utils import export_requirements_from_toml
 
 
 def deploy_to_gcp(
+    gcp_fname: str,
     function_file: str,
     requirements_file: t.Optional[str],
     extra_deps: list[str],
-    labels: dict[str, str],
-    env_vars: dict[str, str],
-    secrets: dict[str, str],
-    market_type: MarketType,
+    labels: dict[str, str] | None,
+    env_vars: dict[str, str] | None,
+    secrets: dict[str, str] | None,
     memory: int,  # in MB
 ) -> str:
     if requirements_file and not os.path.exists(requirements_file):
@@ -37,8 +35,6 @@ def deploy_to_gcp(
     if not os.path.exists(function_file):
         raise ValueError(f"File {function_file} does not exist")
 
-    gcp_fname = DeployableAgent().get_gcloud_fname(market_type=market_type)
-
     # Make a tempdir to store the requirements file and the function
     with tempfile.TemporaryDirectory() as tempdir:
         # Copy function_file to tempdir/main.py

prediction_market_agent_tooling/deploy/gcp/utils.py

@@ -11,9 +11,9 @@ def gcloud_deploy_cmd(
     gcp_function_name: str,
     source: str,
     entry_point: str,
-    labels: dict[str, str],
-    env_vars: dict[str, str],
-    secrets: dict[str, str],
+    labels: dict[str, str] | None,
+    env_vars: dict[str, str] | None,
+    secrets: dict[str, str] | None,
     memory: int,  # in MB
 ) -> str:
     cmd = (
@@ -27,12 +27,15 @@ def gcloud_deploy_cmd(
         f"--memory {memory}MB "
         f"--no-allow-unauthenticated "
     )
-    for k, v in labels.items():
-        cmd += f"--update-labels {k}={v} "
-    for k, v in env_vars.items():
-        cmd += f"--set-env-vars {k}={v} "
-    for k, v in secrets.items():
-        cmd += f"--set-secrets {k}={v} "
+    if labels:
+        for k, v in labels.items():
+            cmd += f"--update-labels {k}={v} "
+    if env_vars:
+        for k, v in env_vars.items():
+            cmd += f"--set-env-vars {k}={v} "
+    if secrets:
+        for k, v in secrets.items():
+            cmd += f"--set-secrets {k}={v} "
 
     return cmd
 

tests/deploy/test_deploy.py

@@ -1,6 +1,6 @@
 import random
 
-from prediction_market_agent_tooling.deploy.agent import DeployableAgent, DeploymentType
+from prediction_market_agent_tooling.deploy.agent import DeployableAgent
 from prediction_market_agent_tooling.markets.data_models import AgentMarket
 from prediction_market_agent_tooling.markets.markets import MarketType
 
@@ -16,10 +16,9 @@ def answer_binary_market(self, market: AgentMarket) -> bool:
             return random.choice([True, False])
 
     agent = DeployableCoinFlipAgent()
-    agent.deploy(
+    agent.deploy_local(
         sleep_time=0.001,
         market_type=MarketType.MANIFOLD,
-        deployment_type=DeploymentType.LOCAL,
         timeout=0.01,
         place_bet=False,
     )