Merge pull request agentos-project#356 from andyk/run_table_in_webapp

Run table in webapp

agentos/agent_run.py

@@ -5,6 +5,8 @@
 from mlflow.entities import RunStatus
 from mlflow.utils.mlflow_tags import MLFLOW_PARENT_RUN_ID, MLFLOW_RUN_NAME
 
+from pcs.component_run import ComponentRun
+from pcs.registry import InMemoryRegistry, Registry
 from pcs.run import Run
 
 _EPISODE_KEY = "episode_count"
@@ -56,60 +58,119 @@ class AgentRun(Run):
     RESTORE_KEY = "restore"
     EVALUATE_KEY = "evaluate"
     RUN_TYPE_TAG = "run_type"
-    AGENT_NAME_KEY = "agent_name"
-    ENV_NAME_KEY = "environment_name"
+    AGENT_ID_KEY = "agent_identifier"
+    ENV_ID_KEY = "environment_identifier"
 
     def __init__(
         self,
-        run_type: str,
-        parent_run: Optional[str] = None,
-        agent_name: Optional[str] = None,
-        environment_name: Optional[str] = None,
+        run_type: str = None,
+        parent_run: Run = None,
+        agent_identifier: Optional[str] = None,
+        environment_identifier: Optional[str] = None,
+        existing_run_id: str = None,
     ) -> None:
         """
         Create a new AgentRun.
 
         :param run_type: must be 'evaluate' or 'learn'
-        :param parent_run: Optionally, specify the identifier of another Run
-            that this run is a sub-run of. Setting this will result in this
-            AgentRun being visually nested under the parent_run in the MLflow
-            UI.
-        :param agent_name: The name of the agent component being evaluated or
-            trained. Defaults to "agent".
-        :param environment_name: The name of the environment component being
-            evaluated or trained. Defaults to "environment".
+        :param parent_run: Optionally, specify another Run that this run is
+            a sub-run of. Setting this will result in this AgentRun being
+            visually nested under the parent_run in the MLflow UI.
+        :param agent_identifier: Identifier of Agent component being evaluated
+            or trained.
+        :param environment_identifier: Identifier of Environment component
+            being evaluated or trained.
+        :param existing_run_id: Optional. If provided, load an existing run
+            from the MLflow backing store. If provided, no other options can be
+            provided.
         """
-        super().__init__()
-        self.parent_run = parent_run
-        self.set_tag(self.IS_AGENT_RUN_TAG, "True")
-        self.episode_data = []
-        self.run_type = run_type
-        self.agent_name = agent_name or "agent"
-        self.environment_name = environment_name or "environment"
-
-        self.set_tag(
-            MLFLOW_RUN_NAME,
-            (
-                f"AgentOS {run_type} with Agent '{self.agent_name}' "
-                f"and Env '{self.environment_name}'"
-            ),
+        if existing_run_id:
+            assert not (
+                run_type
+                or parent_run
+                or agent_identifier
+                or environment_identifier
+            ), (
+                "If 'existing_run_id' is specified, then 'run_type', "
+                "'parent_run', 'agent_identifier', and "
+                "'environment_identifier' must be None."
+            )
+
+            super().__init__(existing_run_id=existing_run_id)
+            parent_run_id = self.data.tags[MLFLOW_PARENT_RUN_ID]
+            self.parent_run = ComponentRun.from_existing_run_id(parent_run_id)
+            self.episode_data = []
+            self.run_type = self.data.tags[self.RUN_TYPE_TAG]
+            self.agent_identifier = self.data.tags[self.AGENT_ID_KEY]
+            self.environment_identifier = self.data.tags[self.ENV_ID_KEY]
+        else:
+            assert agent_identifier and environment_identifier, (
+                "If 'existing_run_id' is not provided, then "
+                "'agent_identifier' and 'environment_identifier' must be."
+            )
+            super().__init__()
+            self.parent_run = parent_run
+            self.set_tag(self.IS_AGENT_RUN_TAG, "True")
+            self.episode_data = []
+            self.run_type = run_type
+            self.agent_identifier = agent_identifier
+            self.environment_identifier = environment_identifier
+
+            self.set_tag(
+                MLFLOW_RUN_NAME,
+                (
+                    f"AgentOS {run_type} with Agent '{self.agent_identifier}' "
+                    f"and Env '{self.environment_identifier}'"
+                ),
+            )
+            if self.parent_run:
+                self.set_tag(MLFLOW_PARENT_RUN_ID, self.parent_run.info.run_id)
+
+            self.log_run_type(self.run_type)
+            self.log_agent_identifier(self.agent_identifier)
+            self.log_environment_identifier(self.environment_identifier)
+
+    @classmethod
+    def evaluate_run(
+        cls,
+        parent_run: Run = None,
+        agent_identifier: Optional[str] = None,
+        environment_identifier: Optional[str] = None,
+        existing_run_id: str = None,
+    ) -> "AgentRun":
+        return cls(
+            run_type=cls.EVALUATE_KEY,
+            parent_run=parent_run,
+            agent_identifier=agent_identifier,
+            environment_identifier=environment_identifier,
+            existing_run_id=existing_run_id,
         )
-        if self.parent_run:
-            self.set_tag(MLFLOW_PARENT_RUN_ID, self.parent_run.info.run_id)
 
-        self.log_run_type(self.run_type)
-        self.log_agent_name(self.agent_name)
-        self.log_environment_name(self.environment_name)
+    @classmethod
+    def learn_run(
+        cls,
+        parent_run: Run = None,
+        agent_identifier: Optional[str] = None,
+        environment_identifier: Optional[str] = None,
+        existing_run_id: str = None,
+    ) -> "AgentRun":
+        return cls(
+            run_type=cls.LEARN_KEY,
+            parent_run=parent_run,
+            agent_identifier=agent_identifier,
+            environment_identifier=environment_identifier,
+            existing_run_id=existing_run_id,
+        )
 
     def log_run_type(self, run_type: str) -> None:
         self.run_type = run_type
         self.set_tag(self.RUN_TYPE_TAG, self.run_type)
 
-    def log_agent_name(self, agent_name: str) -> None:
-        self.log_param(self.AGENT_NAME_KEY, agent_name)
+    def log_agent_identifier(self, agent_identifier: str) -> None:
+        self.set_tag(self.AGENT_ID_KEY, agent_identifier)
 
-    def log_environment_name(self, environment_name: str) -> None:
-        self.log_param(self.ENV_NAME_KEY, environment_name)
+    def log_environment_identifier(self, environment_identifier: str) -> None:
+        self.set_tag(self.ENV_ID_KEY, environment_identifier)
 
     def log_run_metrics(self):
         assert self.episode_data, "No episode data!"
@@ -132,6 +193,7 @@ def print_results(self):
         if not self.episode_data:
             return
         run_stats = self._get_run_stats()
+        print(f"Results for AgentRun {self.identifier}")
         if self.run_type == self.LEARN_KEY:
             print(
                 "\nTraining results over "
@@ -193,6 +255,26 @@ def add_episode_data(self, steps: int, reward: float):
             }
         )
 
+    def to_registry(
+        self,
+        registry: Registry = None,
+        recurse: bool = True,
+        force: bool = False,
+        include_artifacts: bool = False,
+    ) -> Registry:
+        if not registry:
+            registry = InMemoryRegistry()
+        if recurse:
+            self.parent_run.to_registry(
+                registry=registry,
+                recurse=recurse,
+                force=force,
+                include_artifacts=include_artifacts,
+            )
+        return super().to_registry(
+            registry=registry, force=force, include_artifacts=include_artifacts
+        )
+
     def end(
         self,
         status: str = RunStatus.to_string(RunStatus.FINISHED),

agentos/cli.py

@@ -209,7 +209,8 @@ def status(entity_id, registry_file, use_venv):
 @agentos_cmd.command()
 @_arg_optional_entity_id
 def publish_run(entity_id):
-    Run.from_existing_run_id(run_id=entity_id).publish()
+    r = Run.from_existing_run_id(run_id=entity_id)
+    r.to_registry(Registry.from_default())
 
 
 @agentos_cmd.command()
@@ -249,7 +250,7 @@ def freeze(component_name, registry_file, force, use_venv):
 @_option_registry_file
 @_option_force
 @_option_use_venv
-def publish(
+def publish_component(
     component_name: str, registry_file: str, force: bool, use_venv: bool
 ):
     """
@@ -260,8 +261,8 @@ def publish(
     component = Component.from_registry_file(
         registry_file, component_name, use_venv=use_venv
     )
-    frozen_spec = component.to_frozen_registry(force=force).to_spec()
-    Registry.get_default().add_component_spec(frozen_spec)
+    frozen_component = component.to_versioned_component()
+    frozen_component.to_registry(Registry.get_default(), force=force)
 
 
 @agentos_cmd.command()

documentation/demos/demo_script_sb3_example_agent.py

@@ -0,0 +1,23 @@
+from pcs import Component
+from pcs.registry import WebRegistry
+
+random_agent = Component.from_registry_file(
+    "example_agents/sb3_agent/components.yaml", "sb3_agent"
+).to_versioned_component()
+run = random_agent.run_with_arg_set("evaluate", {})
+wr = WebRegistry("http://localhost:8000/api/v1")
+
+# You can publish the ComponentRun to locally running WebRegistry server
+# via the following command:
+# run.to_registry(wr)
+
+# OR, update and uncomment the code below to publish the AgentRun itself.
+# Look at the output printed by the call to `run_with_arg_set()` above,
+# and copy the ID of the AgentRun from the line that looks similar to this:
+# "Results for AgentRun 85ad300d7d86461c938444063de68343"
+# from agentos import AgentRun
+# agent_run_id = XXXXXXXXXXXXXXXXXXXXXXX
+# from pcs.registry import WebRegistry
+# sb3_agent_run = AgentRun.from_existing_run_id(agent_run_id)
+# wr = WebRegistry("http://localhost:8000/api/v1")
+# sb3_agent_run.to_registry(wr)

documentation/demos/index.rst

@@ -11,10 +11,24 @@ AgentOS WebRegistry django server app, which you can start via::
   # from within your agentos directory.
   python web/manage.py startserver
 
-AgentOS demo agent from Github
-==============================
-In this demo, we compose an agent out of components that we fetch from
-the AgentOS github repo, then we run that agent and publish the run to
+AgentOS example SB3 agent from registry file in Github
+======================================================
+In this demo, we create an agent Component from a registry file that
+exists in our local clone of the AgentOS github repo. Then we auto-convert
+the agent Component from an unversioned Component to a versioned Component.
+Then we run that agent and publish the run to a local WebRegistry.
+
+Assuming you have the WebRegistry server running, try the following code:
+
+.. include:: demo_script_sb3_example_agent.py
+   :literal:
+
+AgentOS example random agent from Github manually assembled
+===========================================================
+This demo is similar to the previous one except instead of loading an
+agent from a registry file that exists in the AgentOS Github repo,
+we manually compose an agent out of components that we fetch from
+the AgentOS github repo. Then we run that agent and publish the run to
 a local WebRegistry. Finally, it fetches the run back in from the
 local WebRegistry and re-runs it.
 
@@ -23,8 +37,8 @@ Assuming you have the WebRegistry server running, try the following code:
 .. include:: demo_script_agentos_github_random_agent.py
    :literal:
 
-Stable Baselines 3
-==================
+Component from registry auto-inferred directly from Stable Baselines 3 Github
+=============================================================================
 In this demo, we infer a Registry automatically from the github repo
 of a popular open source RL framework
 (`Stable Baselines 3 <https://github.com/DLR-RM/stable-baselines3>`_, a fork of
@@ -33,11 +47,11 @@ that registry.
 
 Assuming you have the WebRegistry server running, try the following code:
 
-.. include:: demo_script_sb3.py
+.. include:: demo_script_sb3_single_component_inferred.py
    :literal:
 
-PyTorch RL Algos by Ilya Kostrikov
-==================================
+[Work In Progress] PyTorch RL Algos by Ilya Kostrikov
+=====================================================
 **This demo is under construction.**
 
 In this demo, we infer a Registry from the github repo of another popular

documentation/overview.rst

@@ -1,9 +1,10 @@
 ***********************************
-Component System Overview
+Python Component System (PCS) Intro
 ***********************************
 
-PCS is an implementation of what we call the **component programming model**.
-The goal of the component programming model is to provide:
+The Python Component System (PCS) is an implementation of what we call the
+**component programming model**. The goal of the component programming model
+is to provide:
 
 * Reproducibility
 * Experiment Tracking
@@ -60,8 +61,8 @@ More specifically, a Component consists of the following information:
 
 A developer uses a Component by running an **entry point** on it.  An entry
 point is a managed, runnable function on the Component.  A Component may have
-one or more entry points, and an entry point may be passed arguments (via an
-**Argument Set**) when run.  When an entry point is run in the PCS runtime, it
+one or more entry points, and an entry point may be passed arguments (via a
+:py:func:`pcs.argument_set.ArgumentSet`) when run.  When an entry point is run in the PCS runtime, it
 automatically has its inputs and outputs instrumented and recorded for
 reproducibility and sharing purposes.
 

example_agents/acme_dqn/agent.py

@@ -20,8 +20,11 @@ def __init__(self, **kwargs):
         )
 
     def evaluate(self, num_episodes):
-
-        with self.AcmeRun("evaluate", active_component_run(self)) as run:
+        with self.AcmeRun.evaluate_run(
+            parent_run=active_component_run(self),
+            agent_identifier=self.__component__.identifier,
+            environment_identifier=self.environment.__component__.identifier,
+        ) as run:
             num_episodes = int(num_episodes)
             loop = acme.EnvironmentLoop(
                 self.environment,
@@ -32,7 +35,11 @@ def evaluate(self, num_episodes):
             loop.run(num_episodes=num_episodes)
 
     def learn(self, num_episodes):
-        with self.AcmeRun("learn", active_component_run(self)) as run:
+        with self.AcmeRun.learn_run(
+            parent_run=active_component_run(self),
+            agent_identifier=self.__component__.identifier,
+            environment_identifier=self.environment.__component__.identifier,
+        ) as run:
             num_episodes = int(num_episodes)
             loop = acme.EnvironmentLoop(
                 self.environment, self.agent, should_update=True, logger=run

example_agents/acme_r2d2/agent.py

@@ -10,7 +10,11 @@ def __init__(self, *args, **kwargs):
         pass
 
     def evaluate(self, num_episodes):
-        with self.AcmeRun("evaluate", active_component_run(self)) as run:
+        with self.AcmeRun.evaluate_run(
+            parent_run=active_component_run(self),
+            agent_identifier=self.__component__.identifier,
+            environment_identifier=self.environment.__component__.identifier,
+        ) as run:
             loop = acme.EnvironmentLoop(
                 self.environment,
                 self,
@@ -20,7 +24,11 @@ def evaluate(self, num_episodes):
             loop.run(num_episodes=int(num_episodes))
 
     def learn(self, num_episodes):
-        with self.AcmeRun("learn", active_component_run(self)) as run:
+        with self.AcmeRun.learn_run(
+            parent_run=active_component_run(self),
+            agent_identifier=self.__component__.identifier,
+            environment_identifier=self.environment.__component__.identifier,
+        ) as run:
             loop = acme.EnvironmentLoop(
                 self.environment,
                 self,