polyaxon 2.1.0rc9__py3-none-any.whl → 2.6.0__py3-none-any.whl

polyaxon/_local_process/converter/base/main.py

@@ -0,0 +1,74 @@
+from typing import Dict, Iterable, List, Optional
+
+from clipped.utils.lists import to_list
+
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._flow import V1Init, V1Plugins
+from polyaxon._local_process import process_types
+from polyaxon._runner.converter import BaseConverter as _BaseConverter
+from polyaxon.exceptions import PolyaxonConverterError
+
+
+class MainConverter(_BaseConverter):
+    def _get_main_container(
+        self,
+        container_id: str,
+        main_container: process_types.V1Container,
+        plugins: V1Plugins,
+        artifacts_store: Optional[V1Connection],
+        init: Optional[List[V1Init]],
+        connections: Optional[List[str]],
+        connection_by_names: Dict[str, V1Connection],
+        secrets: Optional[Iterable[V1ConnectionResource]],
+        config_maps: Optional[Iterable[V1ConnectionResource]],
+        run_path: Optional[str],
+        kv_env_vars: List[List] = None,
+        ports: List[int] = None,
+    ) -> process_types.V1Container:
+        connections = connections or []
+        connection_by_names = connection_by_names or {}
+        secrets = secrets or []
+        config_maps = config_maps or []
+
+        if artifacts_store and not run_path:
+            raise PolyaxonConverterError("Run path is required for main container.")
+
+        if artifacts_store and (
+            not plugins.collect_artifacts or plugins.mount_artifacts_store
+        ):
+            if artifacts_store.name not in connection_by_names:
+                connection_by_names[artifacts_store.name] = artifacts_store
+            if artifacts_store.name not in connections:
+                connections.append(artifacts_store.name)
+
+        requested_connections = [connection_by_names[c] for c in connections]
+        requested_config_maps = V1Connection.get_requested_resources(
+            resources=config_maps,
+            connections=requested_connections,
+            resource_key="config_map",
+        )
+        requested_secrets = V1Connection.get_requested_resources(
+            resources=secrets, connections=requested_connections, resource_key="secret"
+        )
+
+        # Env vars
+        env = self._get_main_env_vars(
+            plugins=plugins,
+            kv_env_vars=kv_env_vars,
+            artifacts_store_name=artifacts_store.name if artifacts_store else None,
+            connections=requested_connections,
+            secrets=requested_secrets,
+            config_maps=requested_config_maps,
+        )
+
+        # Env from
+        resources = to_list(requested_secrets, check_none=True) + to_list(
+            requested_config_maps, check_none=True
+        )
+        env += self._get_env_from_json_resources(resources=resources)
+
+        return self._patch_container(
+            container=main_container,
+            name=container_id,
+            env=env,
+        )

polyaxon/_local_process/converter/base/mounts.py

@@ -0,0 +1,82 @@
+from typing import List, Optional
+
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._contexts import paths as ctx_paths
+from polyaxon._local_process import process_types
+from polyaxon._runner.converter import BaseConverter
+
+
+class MountsMixin(BaseConverter):
+    @classmethod
+    def _get_mount_from_store(
+        cls,
+        store: V1Connection,
+    ):
+        pass
+
+    @classmethod
+    def _get_mount_from_resource(
+        cls,
+        resource: V1ConnectionResource,
+    ):
+        pass
+
+    @classmethod
+    def _get_volume(
+        cls,
+        mount_path: str,
+        host_path: Optional[str] = None,
+        read_only: Optional[bool] = None,
+    ):
+        pass
+
+    @classmethod
+    def _get_docker_context_mount(cls):
+        pass
+
+    @classmethod
+    def _get_auth_context_mount(
+        cls,
+        read_only: Optional[bool] = None,
+        run_path: Optional[str] = None,
+    ):
+        pass
+
+    @classmethod
+    def _get_artifacts_context_mount(
+        cls,
+        read_only: bool = False,
+        run_path: Optional[str] = None,
+    ):
+        pass
+
+    @classmethod
+    def _get_connections_context_mount(
+        cls,
+        name: str,
+        mount_path: str,
+        run_path: str,
+    ):
+        pass
+
+    @classmethod
+    def _get_shm_context_mount(cls):
+        """
+        Mount a tmpfs volume to /dev/shm.
+        This will set /dev/shm size to half of the RAM of node.
+        By default, /dev/shm is very small, only 64MB.
+        Some experiments will fail due to lack of share memory,
+        such as some experiments running on Pytorch.
+        """
+        pass
+
+    @classmethod
+    def _get_mounts(
+        cls,
+        use_auth_context: bool,
+        use_docker_context: bool,
+        use_shm_context: bool,
+        use_artifacts_context: bool,
+        run_path: Optional[str] = None,
+    ) -> List:
+        return []

polyaxon/_local_process/converter/converters/__init__.py

@@ -0,0 +1,8 @@
+from polyaxon._flow import V1RunKind
+from polyaxon._local_process.converter.converters.job import JobConverter
+from polyaxon._local_process.converter.converters.service import ServiceConverter
+
+CONVERTERS = {
+    V1RunKind.JOB: JobConverter,
+    V1RunKind.SERVICE: ServiceConverter,
+}

polyaxon/_local_process/converter/converters/job.py

@@ -0,0 +1,40 @@
+from typing import Dict, Iterable, List, Optional
+
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._flow import V1CompiledOperation, V1Job, V1Plugins
+from polyaxon._k8s.converter.mixins import JobMixin
+from polyaxon._local_process import process_types
+from polyaxon._local_process.converter.base import BaseConverter
+
+
+class JobConverter(JobMixin, BaseConverter):
+    def get_resource(
+        self,
+        compiled_operation: V1CompiledOperation,
+        artifacts_store: V1Connection,
+        connection_by_names: Dict[str, V1Connection],
+        secrets: Optional[Iterable[V1ConnectionResource]],
+        config_maps: Optional[Iterable[V1ConnectionResource]],
+        default_sa: Optional[str] = None,
+        default_auth: bool = False,
+    ) -> List[process_types.V1Container]:
+        job = compiled_operation.run  # type: V1Job
+        plugins = V1Plugins.get_or_create(
+            config=compiled_operation.plugins, auth=default_auth
+        )
+        kv_env_vars = compiled_operation.get_env_io()
+        return self.get_replica_resource(
+            environment=job.environment,
+            plugins=plugins,
+            volumes=job.volumes,
+            init=job.init,
+            sidecars=job.sidecars,
+            container=job.container,
+            artifacts_store=artifacts_store,
+            connections=job.connections,
+            connection_by_names=connection_by_names,
+            secrets=secrets,
+            config_maps=config_maps,
+            kv_env_vars=kv_env_vars,
+            default_sa=default_sa,
+        )

polyaxon/_local_process/converter/converters/service.py

@@ -0,0 +1,41 @@
+from typing import Dict, Iterable, List, Optional
+
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._flow import V1CompiledOperation, V1Plugins, V1Service
+from polyaxon._local_process import process_types
+from polyaxon._local_process.converter.base import BaseConverter
+from polyaxon._local_process.converter.mixins import ServiceMixin
+
+
+class ServiceConverter(ServiceMixin, BaseConverter):
+    def get_resource(
+        self,
+        compiled_operation: V1CompiledOperation,
+        artifacts_store: V1Connection,
+        connection_by_names: Dict[str, V1Connection],
+        secrets: Optional[Iterable[V1ConnectionResource]],
+        config_maps: Optional[Iterable[V1ConnectionResource]],
+        default_sa: Optional[str] = None,
+        default_auth: bool = False,
+    ) -> List[process_types.V1Container]:
+        service = compiled_operation.run  # type: V1Service
+        plugins = V1Plugins.get_or_create(
+            config=compiled_operation.plugins, auth=default_auth
+        )
+        kv_env_vars = compiled_operation.get_env_io()
+        return self.get_replica_resource(
+            plugins=plugins,
+            environment=service.environment,
+            volumes=service.volumes,
+            init=service.init,
+            sidecars=service.sidecars,
+            container=service.container,
+            artifacts_store=artifacts_store,
+            connections=service.connections,
+            connection_by_names=connection_by_names,
+            secrets=secrets,
+            config_maps=config_maps,
+            kv_env_vars=kv_env_vars,
+            default_sa=default_sa,
+            ports=service.ports,
+        )

polyaxon/_local_process/converter/mixins.py

@@ -0,0 +1,38 @@
+from typing import Dict
+
+from polyaxon._containers.names import MAIN_JOB_CONTAINER
+from polyaxon._flow import V1RunKind
+
+
+class JobMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.JOB
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+class NotifierMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.NOTIFIER
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+class CleanerMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.CLEANER
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+class TunerMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.TUNER
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+class ServiceMixin:
+    K8S_ANNOTATIONS_KIND = V1RunKind.SERVICE
+    MAIN_CONTAINER_ID = MAIN_JOB_CONTAINER
+
+
+MIXIN_MAPPING: Dict = {
+    V1RunKind.JOB: JobMixin,
+    V1RunKind.NOTIFIER: NotifierMixin,
+    V1RunKind.CLEANER: CleanerMixin,
+    V1RunKind.TUNER: TunerMixin,
+    V1RunKind.SERVICE: ServiceMixin,
+}

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,42 @@
+from typing import Dict, List, Optional, Tuple, Union
+
+from clipped.compact.pydantic import PYDANTIC_VERSION, Field
+
+from polyaxon._schemas.base import BaseSchemaModel, RootModel
+
+
+class V1EnvVar(RootModel):
+    if PYDANTIC_VERSION.startswith("2."):
+        root: Union[Tuple[str, str], Dict[str, str]]
+    else:
+        __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] = None
+    command: Optional[List[str]] = None
+    args: Optional[List[str]] = None
+    env: Optional[List[V1EnvVar]] = None
+    working_dir: Optional[str] = Field(alias="workingDir", default=None)
+
+    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/_polyaxonfile/specs/compiled_operation.py

@@ -348,7 +348,7 @@ class CompiledOperationSpecification(BaseSpecification):
             "conditions",
             "skip_on_upstream_skip",
         }
-        patch_keys = patch_keys.intersection(preset.__fields_set__)
+        patch_keys = patch_keys.intersection(preset.model_fields_set)
         patch_data = {k: getattr(preset, k) for k in patch_keys}
         patch_compiled = V1CompiledOperation.construct(**patch_data)
         return config.patch(patch_compiled, strategy=preset.patch_strategy)

polyaxon/_polyaxonfile/specs/libs/parser.py

@@ -93,7 +93,7 @@ class PolyaxonfileParser:
         # Check workflow
         for section in Sections.PARSING_SECTIONS:
             config_section = cls._get_section(config, section)
-            if config_section:
+            if config_section is not None:
                 parsed_data[section] = cls.parse_expression(
                     config_section, parsed_params
                 )

polyaxon/_polyaxonfile/specs/libs/validator.py

@@ -11,7 +11,7 @@ def validate(spec, data):
 
     def validate_keys(section, config, section_data):
         extra_args = [
-            key for key in section_data.keys() if key not in config.__fields__.keys()
+            key for key in section_data.keys() if key not in config.model_fields.keys()
         ]
         if extra_args:
             raise PolyaxonfileError(

polyaxon/_polyaxonfile/specs/operation.py

@@ -107,7 +107,7 @@ class OperationSpecification(BaseSpecification):
             "conditions",
             "skip_on_upstream_skip",
         }
-        patch_keys = patch_keys.intersection(config.__fields_set__)
+        patch_keys = patch_keys.intersection(config.model_fields_set)
         patch_data = {k: getattr(config, k) for k in patch_keys}
         patch_compiled = V1CompiledOperation.construct(contexts=contexts, **patch_data)
 

polyaxon/_polyaxonfile/specs/sections.py

@@ -99,6 +99,14 @@ class Sections:
         CONDITIONS,
         SKIP_ON_UPSTREAM_SKIP,
         PATCH_STRATEGY,
+        "is_approved",
+        "patch_strategy",
+        "is_preset",
+        "hub_ref",
+        "dag_ref",
+        "path_ref",
+        "url_ref",
+        "skip_on_upstream_skip",
     )
 
     REQUIRED_SECTIONS = (VERSION, KIND)

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
@@ -23,16 +24,16 @@ class BaseAsyncAgent(BaseAgent):
     IS_ASYNC = True
 
     async def _enter(self):
+        logger.warning("Agent is starting.")
+        await self.executor.refresh()
         if not self.client._is_managed:
             return self
-        print("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."
@@ -80,13 +81,16 @@ class BaseAsyncAgent(BaseAgent):
 
     async def reconcile(self):
         if (
-            now() - self._last_reconciled_at
+            now() - self._last_data_collected_at
         ).total_seconds() > self.SLEEP_AGENT_DATA_COLLECT_TIME:
+            await self.collect_agent_data()
+        if (
+            now() - self._last_reconciled_at
+        ).total_seconds() < self.SLEEP_AGENT_DATA_RECONCILE_TIME:
             return
 
-        # Collect data
-        await self.collect_agent_data()
-
+        logger.info("Checking cluster state.")
+        self._last_reconciled_at = now()
         # Update reconcile
         namespaces = [settings.AGENT_CONFIG.namespace]
         namespaces += settings.AGENT_CONFIG.additional_namespaces or []
@@ -96,9 +100,19 @@ class BaseAsyncAgent(BaseAgent):
             if _ops:
                 ops += [
                     (
-                        op["metadata"]["name"],
-                        op["metadata"]["labels"]["app.kubernetes.io/instance"],
-                        op["metadata"]["annotations"]["polyaxon.com/run_kind"],
+                        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
@@ -142,7 +156,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()
 

polyaxon/_runner/agent/base_agent.py

@@ -4,7 +4,7 @@ import traceback
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Dict, Optional, Tuple, Type
 
-from clipped.utils.tz import now
+from clipped.utils.tz import get_datetime_from_now, now
 
 from polyaxon import settings
 from polyaxon._auxiliaries import V1PolyaxonInitContainer, V1PolyaxonSidecarContainer
@@ -24,7 +24,8 @@ class BaseAgent:
     HEALTH_FILE = "/tmp/.healthz"
     SLEEP_STOP_TIME = 60 * 5
     SLEEP_ARCHIVED_TIME = 60 * 60
-    SLEEP_AGENT_DATA_COLLECT_TIME = 60 * 30
+    SLEEP_AGENT_DATA_COLLECT_TIME = 60 * 15
+    SLEEP_AGENT_DATA_RECONCILE_TIME = 60 * 5
     IS_ASYNC = False
 
     def __init__(
@@ -38,11 +39,13 @@ class BaseAgent:
             self.max_interval = max(max_interval, 3)
         if not agent_uuid and not owner:
             owner = DEFAULT
+        last_hour = get_datetime_from_now(days=0, hours=1)
         self.executor = None
         self._default_auth = bool(agent_uuid)
-        self._executor_refreshed_at = now()
+        self._executor_refreshed_at = last_hour
         self._graceful_shutdown = False
-        self._last_reconciled_at = now()
+        self._last_data_collected_at = last_hour
+        self._last_reconciled_at = last_hour
         self.client = AgentClient(
             owner=owner, agent_uuid=agent_uuid, is_async=self.IS_ASYNC
         )
@@ -60,10 +63,16 @@ class BaseAgent:
 
     def collect_agent_data(self):
         logger.info("Collecting agent data.")
-        self._last_reconciled_at = now()
-        return self.client.collect_agent_data(
-            namespace=settings.CLIENT_CONFIG.namespace
-        )
+        self._last_data_collected_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:
@@ -126,14 +135,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."