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

polyaxon/_k8s/logging/async_monitor.py

@@ -1,4 +1,5 @@
 import datetime
+import logging
 
 from typing import List, Optional, Tuple
 
@@ -10,6 +11,8 @@ from polyaxon._flow import V1RunKind
 from polyaxon._k8s.manager.async_manager import AsyncK8sManager
 from traceml.logging import V1Log, V1Logs
 
+_logger = logging.getLogger("haupt.k8s.logs")
+
 
 async def handle_container_logs(
     k8s_manager: AsyncK8sManager, pod: V1Pod, container_name: str, **params
@@ -23,8 +26,20 @@ async def handle_container_logs(
             timestamps=True,
             **params,
         )
-    except ApiException:
-        pass
+    except ApiException as e:
+        _logger.warning(
+            "Error collecting logs for %s - container %s: %s",
+            pod.metadata.name,
+            container_name,
+            e,
+        )
+    except Exception as e:
+        _logger.warning(
+            "Unexpected error collecting logs for %s - container %s: %s",
+            pod.metadata.name,
+            container_name,
+            e,
+        )
     if not resp:
         return []
 
@@ -66,7 +81,7 @@ async def query_k8s_operation_logs(
     new_time = now()
     params = {}
     if last_time:
-        since_seconds = (new_time - last_time).total_seconds() - 1
+        since_seconds = (new_time - last_time).total_seconds()
         params["since_seconds"] = int(since_seconds)
     if stream:
         params["tail_lines"] = V1Logs._CHUNK_SIZE
@@ -83,6 +98,11 @@ async def query_k8s_operation_logs(
             **params,
         )
 
+    if logs and last_time:
+        # make sure to filter logs larger than last_time
+        logs = [log for log in logs if log.timestamp > last_time]
+    if logs and logs[-1].timestamp:
+        new_time = logs[-1].timestamp
     return logs, new_time
 
 
@@ -96,32 +116,35 @@ async def collect_agent_service_logs(
         k8s_manager=k8s_manager,
         pod=pod,
         container_name=container.name,
-        tail_lines=V1Logs._CHUNK_SIZE * 3,
+        tail_lines=V1Logs._CHUNK_SIZE,
     )
 
 
 async def query_k8s_pod_logs(
     k8s_manager: AsyncK8sManager,
     pod: V1Pod,
-    last_time: Optional[datetime.datetime],
+    last_time: Optional[datetime.datetime] = None,
     stream: bool = False,
 ) -> Tuple[List[V1Log], Optional[datetime.datetime]]:
     new_time = now()
     params = {}
     if last_time:
-        since_seconds = (new_time - last_time).total_seconds() - 1
+        since_seconds = (new_time - last_time).total_seconds()
         params["since_seconds"] = int(since_seconds)
     if stream:
         params["tail_lines"] = V1Logs._CHUNK_SIZE
 
     logs = await handle_pod_logs(k8s_manager=k8s_manager, pod=pod, **params)
 
-    if logs:
-        last_time = logs[-1].timestamp
-    return logs, last_time
+    if logs and last_time:
+        # make sure to filter logs larger than last_time
+        logs = [log for log in logs if log.timestamp > last_time]
+    if logs and logs[-1].timestamp:
+        new_time = logs[-1].timestamp
+    return logs, new_time
 
 
-async def get_op_pos_and_services(
+async def get_op_pods_and_services(
     k8s_manager: AsyncK8sManager,
     run_uuid: str,
     run_kind: str,
@@ -138,12 +161,40 @@ async def get_op_pos_and_services(
     return pods, services
 
 
+async def get_resource_events(
+    k8s_manager: AsyncK8sManager, resource_type: str, resource_name: str
+):
+    field_selector = (
+        f"involvedObject.kind={resource_type},involvedObject.name={resource_name}"
+    )
+    try:
+        events = await k8s_manager.list_namespaced_events(field_selector=field_selector)
+
+        all_events = []
+        for event in events:
+            event_data = {
+                "reason": event.reason,
+                "message": event.message,
+                "first_timestamp": event.first_timestamp,
+                "last_timestamp": event.last_timestamp,
+                "count": event.count,
+                "type": event.type,
+            }
+            all_events.append(event_data)
+
+        return all_events
+
+    except ApiException as e:
+        print(f"Exception when calling CoreV1Api->list_namespaced_event: {e}")
+        return []
+
+
 async def get_op_spec(
     k8s_manager: AsyncK8sManager,
     run_uuid: str,
     run_kind: str,
 ):
-    pods, services = await get_op_pos_and_services(
+    pods, services = await get_op_pods_and_services(
         k8s_manager=k8s_manager,
         run_uuid=run_uuid,
         run_kind=run_kind,
@@ -153,11 +204,21 @@ async def get_op_spec(
         pods_list[
             pod.metadata.name
         ] = k8s_manager.api_client.sanitize_for_serialization(pod)
+        pods_list[pod.metadata.name]["events"] = await get_resource_events(
+            k8s_manager=k8s_manager,
+            resource_type="Pod",
+            resource_name=pod.metadata.name,
+        )
     services_list = {}
     for service in services or []:
         services_list[
             service.metadata.name
         ] = k8s_manager.api_client.sanitize_for_serialization(service)
+        services_list[service.metadata.name]["events"] = await get_resource_events(
+            k8s_manager=k8s_manager,
+            resource_type="Service",
+            resource_name=service.metadata.name,
+        )
     data = {"pods": pods_list, "services": services_list}
     return data, pods, services
 
@@ -185,11 +246,21 @@ async def get_agent_spec(
         pods_list[
             pod.metadata.name
         ] = k8s_manager.api_client.sanitize_for_serialization(pod)
+        pods_list[pod.metadata.name]["events"] = await get_resource_events(
+            k8s_manager=k8s_manager,
+            resource_type="Pod",
+            resource_name=pod.metadata.name,
+        )
     data = {"pods": pods_list}
     services_list = {}
     for service in services or []:
         services_list[
             service.metadata.name
         ] = k8s_manager.api_client.sanitize_for_serialization(service)
+        services_list[service.metadata.name]["events"] = await get_resource_events(
+            k8s_manager=k8s_manager,
+            resource_type="Service",
+            resource_name=service.metadata.name,
+        )
     data["services"] = services_list
     return data, pods, services

polyaxon/_k8s/manager/async_manager.py

@@ -267,3 +267,18 @@ class AsyncK8sManager(BaseK8sManager):
                 raise PolyaxonK8sError("Connection error: %s" % e) from e
             else:
                 logger.debug("Custom object `{}` was not found".format(name))
+
+    async def list_namespaced_events(
+        self,
+        field_selector: str = None,
+        namespace: str = None,
+        reraise: bool = False,
+        **kwargs
+    ) -> List[client.CoreV1EventList]:
+        return await self._list_namespace_resource(
+            resource_api=self.k8s_api.list_namespaced_event,  # type: ignore[attr-defined]
+            reraise=reraise,
+            namespace=namespace,
+            field_selector=field_selector,
+            **kwargs,
+        )

polyaxon/_k8s/manager/manager.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import List, Optional
 
 from kubernetes import client, config
 from kubernetes.client import Configuration
@@ -806,3 +806,18 @@ class K8sManager(BaseK8sManager):
             pod_id, namespace=namespace or self.namespace
         )
         return is_pod_running(event, container_id)
+
+    def list_namespaced_events(
+        self,
+        field_selector: str = None,
+        namespace: str = None,
+        reraise: bool = False,
+        **kwargs
+    ) -> List[client.CoreV1EventList]:
+        return self._list_namespace_resource(
+            resource_api=self.k8s_api.list_namespaced_event,  # type: ignore[attr-defined]
+            reraise=reraise,
+            namespace=namespace,
+            field_selector=field_selector,
+            **kwargs,
+        )

polyaxon/_local_process/agent.py

@@ -0,0 +1,6 @@
+from polyaxon._local_process.executor import Executor
+from polyaxon._runner.agent.sync_agent import BaseSyncAgent
+
+
+class Agent(BaseSyncAgent):
+    EXECUTOR = Executor

polyaxon/_local_process/converter/__init__.py

@@ -0,0 +1 @@
+

polyaxon/_local_process/converter/base/__init__.py

@@ -0,0 +1 @@
+from polyaxon._local_process.converter.base.base import BaseConverter

polyaxon/_local_process/converter/base/base.py

@@ -0,0 +1,140 @@
+from typing import Dict, Iterable, List, Optional, Union
+
+from polyaxon import settings
+from polyaxon._auxiliaries import V1PolyaxonSidecarContainer
+from polyaxon._connections import V1Connection, V1ConnectionResource
+from polyaxon._flow import V1Environment, V1Init, V1Plugins
+from polyaxon._k8s import k8s_schemas
+from polyaxon._local_process import process_types
+from polyaxon._local_process.converter.base.containers import ContainerMixin
+from polyaxon._local_process.converter.base.env_vars import EnvMixin
+from polyaxon._local_process.converter.base.init import InitConverter
+from polyaxon._local_process.converter.base.main import MainConverter
+from polyaxon._local_process.converter.base.mounts import MountsMixin
+from polyaxon._runner.converter import BaseConverter as _BaseConverter
+from polyaxon._runner.kinds import RunnerKind
+from polyaxon.exceptions import PolyaxonConverterError
+
+
+class BaseConverter(
+    MainConverter, InitConverter, ContainerMixin, EnvMixin, MountsMixin, _BaseConverter
+):
+    RUNNER_KIND = RunnerKind.PROCESS
+
+    @classmethod
+    def _get_sidecar_container(
+        cls,
+        container_id: str,
+        polyaxon_sidecar: V1PolyaxonSidecarContainer,
+        env: List[process_types.V1EnvVar],
+        artifacts_store: V1Connection,
+        plugins: V1Plugins,
+        run_path: Optional[str],
+    ) -> Optional[process_types.V1Container]:
+        return None
+
+    @classmethod
+    def _k8s_to_process_env_var(
+        cls,
+        env_var: List[k8s_schemas.V1EnvVar],
+    ) -> List[process_types.V1EnvVar]:
+        if not env_var:
+            return []
+
+        process_env_var = []
+        for item in env_var:
+            if isinstance(item, dict):
+                try:
+                    item = k8s_schemas.V1EnvVar(**item)
+                except (ValueError, TypeError) as e:
+                    raise PolyaxonConverterError(
+                        f"Could not parse env var value `{item}` from the K8S schema in container section"
+                    ) from e
+            process_env_var.append(cls._get_env_var(name=item.name, value=item.value))
+
+        return process_env_var
+
+    @staticmethod
+    def _new_container(name: str) -> process_types.V1Container:
+        return process_types.V1Container(name=name)
+
+    @classmethod
+    def _ensure_container(
+        cls,
+        container: Union[k8s_schemas.V1Container, process_types.V1Container],
+        volumes: List[k8s_schemas.V1Volume],
+    ) -> process_types.V1Container:
+        if not isinstance(container, k8s_schemas.V1Container):
+            return container
+        return process_types.V1Container(
+            name=container.name,
+            command=container.command,
+            args=container.args,
+            env=cls._k8s_to_process_env_var(container.env),
+            working_dir=container.working_dir,
+        )
+
+    def get_replica_resource(
+        self,
+        environment: V1Environment,
+        plugins: V1Plugins,
+        volumes: List[k8s_schemas.V1Volume],
+        init: List[V1Init],
+        sidecars: List[k8s_schemas.V1Container],
+        container: k8s_schemas.V1Container,
+        artifacts_store: V1Connection,
+        connections: List[str],
+        connection_by_names: Dict[str, V1Connection],
+        secrets: Optional[Iterable[V1ConnectionResource]],
+        config_maps: Optional[Iterable[V1ConnectionResource]],
+        kv_env_vars: List[List],
+        default_sa: Optional[str] = None,
+        ports: List[int] = None,
+    ) -> List[process_types.V1Container]:
+        volumes = volumes or []
+        init = init or []
+        sidecars = sidecars or []
+        connections = connections or []
+        environment = environment or V1Environment()
+        environment.service_account_name = (
+            environment.service_account_name
+            or default_sa
+            or settings.AGENT_CONFIG.runs_sa
+        )
+
+        init_connections = self.filter_connections_from_init(init=init)
+
+        init_containers = self.get_init_containers(
+            polyaxon_init=self.polyaxon_init,
+            plugins=plugins,
+            artifacts_store=artifacts_store,
+            init_connections=init_connections,
+            init_containers=self.filter_containers_from_init(init=init),
+            connection_by_names=connection_by_names,
+            log_level=plugins.log_level,
+            volumes=volumes,
+        )
+
+        sidecar_containers = self.get_sidecar_containers(
+            polyaxon_sidecar=self.polyaxon_sidecar,
+            plugins=plugins,
+            artifacts_store=artifacts_store,
+            sidecar_containers=sidecars,
+            log_level=plugins.log_level,
+            volumes=volumes,
+        )
+
+        main_container = self.get_main_container(
+            main_container=self._ensure_container(container, volumes=volumes),
+            plugins=plugins,
+            artifacts_store=artifacts_store,
+            connections=connections,
+            init_connections=init_connections,
+            connection_by_names=connection_by_names,
+            secrets=secrets,
+            config_maps=config_maps,
+            ports=ports,
+            kv_env_vars=kv_env_vars,
+        )
+
+        return init_containers + sidecar_containers + [main_container]

polyaxon/_local_process/converter/base/containers.py

@@ -0,0 +1,66 @@
+from typing import Any, Dict, List, Optional
+
+from clipped.utils.lists import to_list
+from clipped.utils.sanitizers import sanitize_value
+
+from polyaxon._containers.names import sanitize_container_name
+from polyaxon._local_process import process_types
+from polyaxon._runner.converter import BaseConverter
+from polyaxon._runner.converter.common.containers import sanitize_container_command_args
+
+
+class ContainerMixin(BaseConverter):
+    @classmethod
+    def _patch_container(
+        cls,
+        container: process_types.V1Container,
+        name: Optional[str] = None,
+        command: Optional[List[str]] = None,
+        args: Optional[List[str]] = None,
+        env: Optional[List[process_types.V1EnvVar]] = None,
+        env_from: Optional[List[Any]] = None,
+        **kwargs,
+    ) -> process_types.V1Container:
+        container.name = sanitize_container_name(name or container.name)
+        container.env = to_list(container.env, check_none=True) + to_list(
+            env, check_none=True
+        )
+        if not any([container.command, container.args]):
+            container.command = command
+            container.args = args
+
+        return cls._sanitize_container(container)
+
+    @staticmethod
+    def _sanitize_container_env(
+        env: List[process_types.V1EnvVar],
+    ) -> List[process_types.V1EnvVar]:
+        def sanitize_env_dict(d: Dict):
+            return process_types.V1EnvVar.make(
+                {d_k: sanitize_value(d_v, handle_dict=False) for d_k, d_v in d.items()}
+            )
+
+        results = []
+        for e in env or []:
+            if isinstance(e, dict):
+                e = sanitize_env_dict(e)
+                results.append(e)
+            elif isinstance(e, tuple):
+                if e[1] is not None:
+                    e = process_types.V1EnvVar.make(
+                        (e[0], sanitize_value(e[1], handle_dict=False))
+                    )
+                results.append(e)
+            elif isinstance(e, process_types.V1EnvVar):
+                results.append(e)
+
+        return results
+
+    @classmethod
+    def _sanitize_container(
+        cls,
+        container: process_types.V1Container,
+    ) -> process_types.V1Container:
+        container = sanitize_container_command_args(container)
+        container.env = cls._sanitize_container_env(container.env)
+        return container