polyaxon 2.0.0rc49__py3-none-any.whl → 2.4.0rc1__py3-none-any.whl

polyaxon/_local_process/executor.py

@@ -0,0 +1,132 @@
+import os
+import signal
+import subprocess
+
+from typing import Dict, List
+
+from polyaxon._deploy.operators.cmd_operator import CmdOperator
+from polyaxon._deploy.operators.conda import CondaOperator
+from polyaxon._local_process import process_types
+from polyaxon._local_process.converter.converters import CONVERTERS
+from polyaxon._local_process.converter.mixins import MIXIN_MAPPING
+from polyaxon._runner.executor import BaseExecutor
+from polyaxon._runner.kinds import RunnerKind
+from polyaxon._schemas.lifecycle import V1Statuses
+from polyaxon.exceptions import PolyaxonAgentError
+from polyaxon.logger import logger
+
+
+class Executor(BaseExecutor):
+    MIXIN_MAPPING = MIXIN_MAPPING
+    CONVERTERS = CONVERTERS
+    RUNNER_KIND = RunnerKind.PROCESS
+
+    def __init__(self, conda_env: str = None, venv: str = None):
+        super().__init__()
+        self._ops = {}
+        self._conda_env = conda_env
+        self._venv = venv
+
+    def _get_manager(self):
+        if self._conda_env:
+            return CondaOperator()
+        return CmdOperator()
+
+    def _check_conda(self):
+        if not self.manager.check():
+            raise logger.error("Conda is required to run this command.")
+
+        envs = self.manager.execute(["env", "list", "--json"], is_json=True)
+        env_names = [os.path.basename(env) for env in envs["envs"]]
+        if self._conda_env not in env_names:
+            raise logger.error(
+                "Conda env `{}` is not installed.".format(self._conda_env),
+                sys_exit=True,
+            )
+
+    def _run_in_conda(self, cmd_bash, cmd_args):
+        cmd_args = ["source activate {}".format(self._conda_env)] + cmd_args
+        subprocess.Popen(cmd_bash + [" && ".join(cmd_args)], close_fds=True)
+
+    def _get_op_proc(self, run_uuid: str) -> List[subprocess.Popen]:
+        return self._ops.get(run_uuid)
+
+    def create(
+        self,
+        run_uuid: str,
+        run_kind: str,
+        resource: List[process_types.V1Container],
+        namespace: str = None,
+    ) -> Dict:
+        logger.info(f"[Executor] Starting operation {run_uuid} {run_kind}.")
+        self._ops[run_uuid] = []
+        for task in resource:
+            logger.info(
+                f"[Executor] Starting task container {task.name} {task.image} ."
+            )
+            proc = self.manager.execute(
+                task.get_cmd_args(), env=os.environ, output_only=False
+            )
+            self._ops[run_uuid].append(proc)
+            proc.wait()
+            task_status = self._get_task_status(proc)
+            message = f"Task container {task.name} {task.image} with id {proc.pid} {task_status}"
+            if task_status == V1Statuses.SUCCEEDED:
+                logger.info(f"[Executor] {message}")
+            else:
+                logger.warning(f"[Executor] {message}")
+                self._clean_temp_execution_path(run_uuid)
+                return {
+                    "status": V1Statuses.FAILED,
+                    "tasks": self._ops[run_uuid],
+                    "message": message,
+                }
+        self._clean_temp_execution_path(run_uuid)
+        return {"status": V1Statuses.SUCCEEDED, "tasks": self._ops[run_uuid]}
+
+    def apply(
+        self, run_uuid: str, run_kind: str, resource: Dict, namespace: str = None
+    ) -> Dict:
+        raise PolyaxonAgentError(
+            "Docker executor does not support apply method.\n"
+            "Run: <kind: {}, uuid: {}>".format(run_kind, run_uuid)
+        )
+
+    def stop(self, run_uuid: str, run_kind: str, namespace: str = None):
+        proc = self._get_op_proc(run_uuid)
+        if proc.poll() is None:
+            # Kill the process tree rooted at the child if it's the leader of its own process
+            # group, otherwise just kill the child
+            try:
+                if proc.pid == os.getpgid(proc.pid):
+                    os.killpg(proc.pid, signal.SIGTERM)
+                else:
+                    proc.terminate()
+            except OSError:
+                # The child process may have exited before we attempted to terminate it, so we
+                # ignore OSErrors raised during child process termination
+                _msg = f"Failed to terminate operation {run_kind} {run_uuid} child process PID {proc.pid}"
+                logger.debug(_msg)
+            proc.wait()
+
+    def clean(self, run_uuid: str, run_kind: str, namespace: str = None):
+        return self.apply(
+            run_uuid=run_uuid,
+            run_kind=run_kind,
+            resource={"metadata": {"finalizers": None}},
+        )
+
+    def _get_task_status(self, proc) -> V1Statuses:
+        exit_code = proc.poll()
+        if exit_code is None:
+            return V1Statuses.RUNNING
+        if exit_code == 0:
+            return V1Statuses.SUCCEEDED
+        return V1Statuses.FAILED
+
+    def get(self, run_uuid: str, run_kind: str, namespace: str = None) -> V1Statuses:
+        procs = self._get_op_proc(run_uuid)
+        return self._get_task_status(procs[-1])
+
+    def list_ops(self, namespace: str = None):
+        return []

polyaxon/_local_process/process_types.py

@@ -0,0 +1,39 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+from clipped.compact.pydantic import Field
+
+from polyaxon._schemas.base import BaseSchemaModel
+
+
+class V1EnvVar(BaseSchemaModel):
+    __root__: Union[Tuple[str, str], Dict[str, str]]
+
+    def to_cmd(self):
+        if isinstance(self.__root__, tuple):
+            value = self.__root__
+        else:
+            value = self.__root__.items()
+        return [f"{value[0]}={value[1]}"]
+
+
+class V1Container(BaseSchemaModel):
+    name: Optional[str]
+    command: Optional[List[str]]
+    args: Optional[List[str]]
+    env: Optional[List[V1EnvVar]]
+    working_dir: Optional[str] = Field(alias="workingDir")
+
+    def get_cmd_args(self):
+        cmd_args = ["run", "--rm"]
+        for env in self.env:
+            cmd_args += ["-e"] + env.to_cmd()
+        if self.working_dir:
+            cmd_args += ["-w", self.working_dir]
+        if self.command:
+            cmd_args += ["--entrypoint", self.command[0]]
+        cmd_args += [self.image]
+        if self.command:
+            cmd_args += self.command[1:]
+        if self.args:
+            cmd_args += self.args
+        return cmd_args

polyaxon/_managers/agent.py

@@ -13,7 +13,9 @@ class AgentConfigManager(ConfigManager):
 
     VISIBILITY = ConfigManager.Visibility.GLOBAL
     CONFIG_FILE_NAME = ".agent"
+    ALTERNATE_CONFIG_FILE_NAME = ".sandbox"
     CONFIG: Type[AgentConfig] = AgentConfig
+    PERSIST_FORMAT = "yaml"
 
     @classmethod
     def get_config_or_default(cls) -> AgentConfig:

polyaxon/_managers/home.py

@@ -13,6 +13,7 @@ class HomeConfigManager(ConfigManager):
     VISIBILITY = ConfigManager.Visibility.GLOBAL
     CONFIG_FILE_NAME = ".home"
     CONFIG: Type[HomeConfig] = HomeConfig
+    PERSIST_FORMAT = "yaml"
 
     @classmethod
     def get_config_defaults(cls) -> Dict[str, str]:
@@ -23,7 +24,7 @@ class HomeConfigManager(ConfigManager):
         glob_path = cls.get_global_config_path()
         home_config = cls._CONFIG_READER.read_configs(
             [
-                ConfigSpec(glob_path, config_type=".json", check_if_exists=False),
+                ConfigSpec(glob_path, config_type=".yaml", check_if_exists=False),
                 os.environ,
                 {"dummy": "dummy"},
             ]

polyaxon/_operations/tuner.py

@@ -31,6 +31,7 @@ def get_tuner(
 
     return V1Operation(
         queue=tuner.queue,
+        namespace=tuner.namespace,
         joins=[join],
         params=params,
         hub_ref=tuner.hub_ref,

polyaxon/_polyaxonfile/check.py

@@ -80,6 +80,7 @@ def check_polyaxonfile(
     matrix: Optional[Union[Dict, V1Matrix]] = None,
     presets: Optional[List[str]] = None,
     queue: Optional[str] = None,
+    namespace: Optional[str] = None,
     nocache: Optional[bool] = None,
     cache: Optional[Union[int, str, bool]] = None,
     verbose: bool = True,
@@ -177,6 +178,7 @@ def check_polyaxonfile(
                 matrix=matrix,
                 presets=presets,
                 queue=queue,
+                namespace=namespace,
                 nocache=nocache,
                 cache=cache,
                 approved=approved,

polyaxon/_polyaxonfile/manager/operations.py

@@ -29,6 +29,7 @@ def get_op_specification(
     matrix: Optional[Union[Dict, V1Matrix]] = None,
     presets: Optional[List[str]] = None,
     queue: Optional[str] = None,
+    namespace: Optional[str] = None,
     nocache: Optional[bool] = None,
     cache: Optional[Union[int, str, bool]] = None,
     approved: Optional[Union[int, str, bool]] = None,
@@ -70,6 +71,8 @@ def get_op_specification(
         # Check only
         get_queue_info(queue)
         op_data["queue"] = queue
+    if namespace:
+        op_data["namespace"] = namespace
     if cache is not None:
         op_data["cache"] = {"disable": not to_bool(cache)}
     if nocache:

polyaxon/_polyaxonfile/manager/workflows.py

@@ -20,6 +20,7 @@ def get_op_from_schedule(
     op_spec.skip_on_upstream_skip = None
     op_spec.cache = compiled_operation.cache
     op_spec.queue = compiled_operation.queue
+    op_spec.namespace = compiled_operation.namespace
     op_spec.component.inputs = compiled_operation.inputs
     op_spec.component.outputs = compiled_operation.outputs
     op_spec.component.run = compiled_operation.run
@@ -59,6 +60,7 @@ def get_ops_from_suggestions(
         op_spec.skip_on_upstream_skip = None
         op_spec.cache = compiled_operation.cache
         op_spec.queue = compiled_operation.queue
+        op_spec.namespace = compiled_operation.namespace
         op_spec.params = params
         op_spec.component.inputs = compiled_operation.inputs
         op_spec.component.outputs = compiled_operation.outputs

polyaxon/_polyaxonfile/specs/compiled_operation.py

@@ -333,6 +333,7 @@ class CompiledOperationSpecification(BaseSpecification):
             "is_approved",
             "presets",
             "queue",
+            "namespace",
             "cache",
             "build",
             "hooks",

polyaxon/_polyaxonfile/specs/operation.py

@@ -92,6 +92,7 @@ class OperationSpecification(BaseSpecification):
             "is_approved",
             "presets",
             "queue",
+            "namespace",
             "cache",
             "build",
             "hooks",

polyaxon/_polyaxonfile/specs/sections.py

@@ -13,6 +13,7 @@ class Sections:
     QUEUE = "queue"
     CACHE = "cache"
     PLUGINS = "plugins"
+    NAMESPACE = "namespace"
     BUILD = "build"
     HOOKS = "hooks"
     EVENTS = "events"
@@ -55,6 +56,7 @@ class Sections:
         CACHE,
         QUEUE,
         PLUGINS,
+        NAMESPACE,
         BUILD,
         HOOKS,
         EVENTS,
@@ -89,6 +91,7 @@ class Sections:
         CACHE,
         CONNECTIONS,
         PLUGINS,
+        NAMESPACE,
         TERMINATION,
         SCHEDULE,
         DEPENDENCIES,

polyaxon/_pql/manager.py

@@ -13,7 +13,7 @@ class PQLManager:
     FIELDS_PROXY = {}
     FIELDS_TRANS = {}
     FIELDS_ORDERING = None
-    FIELDS_ORDERING_PROXY = None
+    FIELDS_ORDERING_PROXY = None  # Do not set a field on both field and proxy
     FIELDS_DEFAULT_ORDERING = None
     FIELDS_DISTINCT = None
     CHECK_ALIVE = True

polyaxon/_runner/agent/async_agent.py

@@ -14,6 +14,7 @@ from polyaxon._env_vars.getters import get_run_info
 from polyaxon._runner.agent.base_agent import BaseAgent
 from polyaxon._sdk.schemas.v1_agent import V1Agent
 from polyaxon._sdk.schemas.v1_agent_state_response import V1AgentStateResponse
+from polyaxon._utils.fqn_utils import get_run_instance
 from polyaxon.exceptions import ApiException as SDKApiException
 from polyaxon.exceptions import PolyaxonAgentError, PolyaxonConverterError
 from polyaxon.logger import logger
@@ -25,14 +26,14 @@ class BaseAsyncAgent(BaseAgent):
     async def _enter(self):
         if not self.client._is_managed:
             return self
-        print("Agent is starting.")
+        logger.warning("Agent is starting.")
         await self.executor.refresh()
         try:
             agent = await self.client.get_info()
             self._check_status(agent)
             await self.sync()
             await self.client.log_agent_running()
-            print("Agent is running.")
+            logger.warning("Agent is running.")
             return self
         except (ApiException, SDKApiException, HTTPError) as e:
             message = "Could not start the agent."
@@ -78,6 +79,49 @@ class BaseAsyncAgent(BaseAgent):
             ),
         )
 
+    async def reconcile(self):
+        if (
+            now() - self._last_reconciled_at
+        ).total_seconds() > self.SLEEP_AGENT_DATA_COLLECT_TIME:
+            return
+
+        # Collect data
+        await self.collect_agent_data()
+
+        # Update reconcile
+        namespaces = [settings.AGENT_CONFIG.namespace]
+        namespaces += settings.AGENT_CONFIG.additional_namespaces or []
+        ops = []
+        for namespace in namespaces:
+            _ops = await self.executor.list_ops(namespace=namespace)
+            if _ops:
+                ops += [
+                    (
+                        get_run_instance(
+                            owner=op["metadata"]["annotations"][
+                                "operation.polyaxon.com/owner"
+                            ],
+                            project=op["metadata"]["annotations"][
+                                "operation.polyaxon.com/project"
+                            ],
+                            run_uuid=op["metadata"]["labels"][
+                                "app.kubernetes.io/instance"
+                            ],
+                        ),
+                        op["metadata"]["annotations"]["operation.polyaxon.com/kind"],
+                        op["metadata"]["annotations"]["operation.polyaxon.com/name"],
+                        namespace,
+                    )
+                    for op in _ops
+                ]
+        if not ops:
+            return None
+
+        logger.info("Reconcile agent.")
+        return await self.client.reconcile_agent(
+            reconcile={"ops": ops},
+        )
+
     async def start(self):
         try:
             async with async_exit_context() as exit_event:
@@ -91,7 +135,9 @@ class BaseAsyncAgent(BaseAgent):
                     except asyncio.TimeoutError:
                         index += 1
                         await self.refresh_executor()
-                        if not self._default_auth:
+                        if self._default_auth:
+                            await self.reconcile()
+                        else:
                             await self.cron()
                         agent_state = await self.process()
                         if not agent_state:
@@ -107,7 +153,7 @@ class BaseAsyncAgent(BaseAgent):
                         timeout = get_wait(index, max_interval=self.max_interval)
                         logger.info("Sleeping for {} seconds".format(timeout))
         except Exception as e:
-            print(e)
+            logger.warning("Agent failed to start: {}".format(repr(e)))
         finally:
             self.end()
 
@@ -118,7 +164,7 @@ class BaseAsyncAgent(BaseAgent):
                 self.sync_compatible_updates(agent_state.compatible_updates)
 
             if agent_state:
-                logger.info("Starting runs submission process.")
+                logger.info("Checking agent state.")
             else:
                 logger.info("No state was found.")
                 return V1AgentStateResponse.construct()
@@ -185,7 +231,7 @@ class BaseAsyncAgent(BaseAgent):
             )
         return None
 
-    async def submit_run(self, run_data: Tuple[str, str, str, str]):
+    async def submit_run(self, run_data: Tuple[str, str, str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
         resource = await self.prepare_run_resource(
             owner_name=run_owner,
@@ -197,9 +243,13 @@ class BaseAsyncAgent(BaseAgent):
         if not resource:
             return
 
+        namespace = None if len(run_data) < 5 else run_data[4]
         try:
             await self.executor.create(
-                run_uuid=run_uuid, run_kind=run_data[1], resource=resource
+                run_uuid=run_uuid,
+                run_kind=run_data[1],
+                resource=resource,
+                namespace=namespace,
             )
         except ApiException as e:
             if e.status == 409:
@@ -222,7 +272,7 @@ class BaseAsyncAgent(BaseAgent):
             )
 
     async def make_and_create_run(
-        self, run_data: Tuple[str, str, str, str], default_auth: bool = False
+        self, run_data: Tuple[str, str, str, str, str], default_auth: bool = False
     ):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
         resource = await self.make_run_resource(
@@ -236,9 +286,14 @@ class BaseAsyncAgent(BaseAgent):
         if not resource:
             return
 
+        namepsace = None if len(run_data) < 5 else run_data[4]
+
         try:
             await self.executor.create(
-                run_uuid=run_uuid, run_kind=run_data[1], resource=resource
+                run_uuid=run_uuid,
+                run_kind=run_data[1],
+                resource=resource,
+                namespace=namepsace,
             )
         except ApiException as e:
             if e.status == 409:
@@ -252,7 +307,7 @@ class BaseAsyncAgent(BaseAgent):
                 )
             )
 
-    async def apply_run(self, run_data: Tuple[str, str, str, str]):
+    async def apply_run(self, run_data: Tuple[str, str, str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
         resource = await self.prepare_run_resource(
             owner_name=run_owner,
@@ -264,9 +319,14 @@ class BaseAsyncAgent(BaseAgent):
         if not resource:
             return
 
+        namespace = None if len(run_data) < 5 else run_data[4]
+
         try:
             await self.executor.apply(
-                run_uuid=run_uuid, run_kind=run_data[1], resource=resource
+                run_uuid=run_uuid,
+                run_kind=run_data[1],
+                resource=resource,
+                namespace=namespace,
             )
             await self.client.log_run_running(
                 run_owner=run_owner, run_project=run_project, run_uuid=run_uuid
@@ -275,12 +335,17 @@ class BaseAsyncAgent(BaseAgent):
             await self.client.log_run_failed(
                 run_owner=run_owner, run_project=run_project, run_uuid=run_uuid, exc=e
             )
-            await self.clean_run(run_uuid=run_uuid, run_kind=run_data[1])
+            await self.clean_run(
+                run_uuid=run_uuid, run_kind=run_data[1], namespace=namespace
+            )
 
-    async def check_run(self, run_data: Tuple[str, str]):
+    async def check_run(self, run_data: Tuple[str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
+        namespace = None if len(run_data) < 3 else run_data[2]
         try:
-            await self.executor.get(run_uuid=run_uuid, run_kind=run_data[1])
+            await self.executor.get(
+                run_uuid=run_uuid, run_kind=run_data[1], namespace=namespace
+            )
         except ApiException as e:
             if e.status == 404:
                 logger.info(
@@ -290,10 +355,13 @@ class BaseAsyncAgent(BaseAgent):
                     run_owner=run_owner, run_project=run_project, run_uuid=run_uuid
                 )
 
-    async def stop_run(self, run_data: Tuple[str, str]):
+    async def stop_run(self, run_data: Tuple[str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
+        namespace = None if len(run_data) < 3 else run_data[2]
         try:
-            await self.executor.stop(run_uuid=run_uuid, run_kind=run_data[1])
+            await self.executor.stop(
+                run_uuid=run_uuid, run_kind=run_data[1], namespace=namespace
+            )
         except ApiException as e:
             if e.status == 404:
                 logger.info("Run does not exist anymore, it could have been stopped.")
@@ -309,16 +377,24 @@ class BaseAsyncAgent(BaseAgent):
                 message="Agent failed stopping run.\n",
             )
 
-    async def delete_run(self, run_data: Tuple[str, str, str, str]):
+    async def delete_run(self, run_data: Tuple[str, str, str, str, str]):
         run_owner, run_project, run_uuid = get_run_info(run_instance=run_data[0])
-        await self.clean_run(run_uuid=run_uuid, run_kind=run_data[1])
+        namespace = None if len(run_data) < 5 else run_data[4]
         if run_data[3]:
             await self.make_and_create_run(run_data)
+        else:
+            await self.clean_run(
+                run_uuid=run_uuid, run_kind=run_data[1], namespace=namespace
+            )
 
-    async def clean_run(self, run_uuid: str, run_kind: str):
+    async def clean_run(self, run_uuid: str, run_kind: str, namespace: str = None):
         try:
-            await self.executor.clean(run_uuid=run_uuid, run_kind=run_kind)
-            await self.executor.stop(run_uuid=run_uuid, run_kind=run_kind)
+            await self.executor.clean(
+                run_uuid=run_uuid, run_kind=run_kind, namespace=namespace
+            )
+            await self.executor.stop(
+                run_uuid=run_uuid, run_kind=run_kind, namespace=namespace
+            )
         except ApiException as e:
             if e.status == 404:
                 logger.info("Run does not exist.")

polyaxon/_runner/agent/base_agent.py

@@ -24,6 +24,7 @@ class BaseAgent:
     HEALTH_FILE = "/tmp/.healthz"
     SLEEP_STOP_TIME = 60 * 5
     SLEEP_ARCHIVED_TIME = 60 * 60
+    SLEEP_AGENT_DATA_COLLECT_TIME = 60 * 30
     IS_ASYNC = False
 
     def __init__(
@@ -41,6 +42,7 @@ class BaseAgent:
         self._default_auth = bool(agent_uuid)
         self._executor_refreshed_at = now()
         self._graceful_shutdown = False
+        self._last_reconciled_at = now()
         self.client = AgentClient(
             owner=owner, agent_uuid=agent_uuid, is_async=self.IS_ASYNC
         )
@@ -50,9 +52,25 @@ class BaseAgent:
     def sync(self):
         raise NotImplementedError
 
+    def reconcile(self):
+        raise NotImplementedError
+
     def cron(self):
         return self.client.cron_agent()
 
+    def collect_agent_data(self):
+        logger.info("Collecting agent data.")
+        self._last_reconciled_at = now()
+        try:
+            return self.client.collect_agent_data(
+                namespace=settings.CLIENT_CONFIG.namespace
+            )
+        except Exception as e:
+            logger.warning(
+                "Agent failed to collect agent data: {}\n"
+                "Retrying ...".format(repr(e))
+            )
+
     def sync_compatible_updates(self, compatible_updates: Dict):
         if compatible_updates and settings.AGENT_CONFIG:
             init = compatible_updates.get("init")
@@ -114,14 +132,14 @@ class BaseAgent:
 
     def _check_status(self, agent_state):
         if agent_state.status == V1Statuses.STOPPED:
-            print(
+            logger.warning(
                 "Agent has been stopped from the platform,"
                 "but the deployment is still running."
                 "Please either set the agent to starting or teardown the agent deployment."
             )
             return self.end(sleep=self.SLEEP_STOP_TIME)
         elif agent_state.live_state < LiveState.LIVE:
-            print(
+            logger.warning(
                 "Agent has been archived from the platform,"
                 "but the deployment is still running."
                 "Please either restore the agent or teardown the agent deployment."
@@ -173,25 +191,25 @@ class BaseAgent:
     ) -> Optional[Any]:
         raise NotImplementedError
 
-    def submit_run(self, run_data: Tuple[str, str, str, str]):
+    def submit_run(self, run_data: Tuple[str, str, str, str, str]):
         raise NotImplementedError
 
     def make_and_create_run(
-        self, run_data: Tuple[str, str, str, str], default_auth: bool = False
+        self, run_data: Tuple[str, str, str, str, str], default_auth: bool = False
     ):
         raise NotImplementedError
 
-    def apply_run(self, run_data: Tuple[str, str, str, str]):
+    def apply_run(self, run_data: Tuple[str, str, str, str, str]):
         raise NotImplementedError
 
-    def check_run(self, run_data: Tuple[str, str]):
+    def check_run(self, run_data: Tuple[str, str, str]):
         raise NotImplementedError
 
-    def stop_run(self, run_data: Tuple[str, str]):
+    def stop_run(self, run_data: Tuple[str, str, str]):
         raise NotImplementedError
 
-    def delete_run(self, run_data: Tuple[str, str, str, str]):
+    def delete_run(self, run_data: Tuple[str, str, str, str, str]):
         raise NotImplementedError
 
-    def clean_run(self, run_uuid: str, run_kind: str):
+    def clean_run(self, run_uuid: str, run_kind: str, namespace: str = None):
         raise NotImplementedError

polyaxon/_runner/agent/client.py

@@ -1,6 +1,6 @@
 import traceback
 
-from typing import Optional
+from typing import Dict, Optional
 
 from polyaxon._schemas.lifecycle import V1StatusCondition, V1Statuses
 from polyaxon.client import PolyaxonClient, V1Agent, V1AgentStateResponse
@@ -65,6 +65,20 @@ class AgentClient:
     def cron_agent(self):
         return self.client.agents_v1.cron_agent(owner=self.owner, _request_timeout=10)
 
+    def collect_agent_data(self, namespace: str):
+        return self.client.internal_agents_v1.collect_agent_data(
+            owner=self.owner,
+            uuid=self.agent_uuid,
+            namespace=namespace,
+        )
+
+    def reconcile_agent(self, reconcile: Dict):
+        return self.client.agents_v1.reconcile_agent(
+            owner=self.owner,
+            uuid=self.agent_uuid,
+            body={"reconcile": reconcile},
+        )
+
     def log_agent_running(self):
         return self.log_agent_status(status=V1Statuses.RUNNING, reason="AgentLogger")