dagster-k8s 0.24.12__py3-none-any.whl → 0.25.0__py3-none-any.whl

dagster_k8s/pipes.py

@@ -1,8 +1,13 @@
+import logging
 import os
 import random
 import re
 import string
+import threading
+import time
+from collections.abc import Callable, Generator
 from contextlib import contextmanager
+from datetime import datetime
 from pathlib import Path
 from typing import Any, Iterator, Mapping, Optional, Sequence, Set, Union
 
@@ -14,6 +19,7 @@ from dagster import (
 from dagster._annotations import public
 from dagster._core.definitions.resource_annotation import TreatAsResourceParam
 from dagster._core.errors import DagsterInvariantViolationError
+from dagster._core.execution.context.asset_execution_context import AssetExecutionContext
 from dagster._core.pipes.client import (
     PipesClient,
     PipesClientCompletedInvocation,
@@ -22,6 +28,7 @@ from dagster._core.pipes.client import (
     PipesParams,
 )
 from dagster._core.pipes.context import PipesMessageHandler
+from dagster._core.pipes.merge_streams import LogItem, merge_streams
 from dagster._core.pipes.utils import (
     PipesEnvContextInjector,
     extract_message_or_forward_to_stdout,
@@ -43,6 +50,10 @@ from dagster_k8s.client import (
 from dagster_k8s.models import k8s_model_from_dict, k8s_snake_case_dict
 from dagster_k8s.utils import get_common_labels
 
+INIT_WAIT_TIMEOUT_FOR_READY = 1800.0  # 30mins
+INIT_WAIT_TIMEOUT_FOR_TERMINATE = 10.0  # 10s
+WAIT_TIMEOUT_FOR_READY = 18000.0  # 5hrs
+
 
 def get_pod_name(run_id: str, op_name: str):
     clean_op_name = re.sub("[^a-z0-9-]", "", op_name.lower().replace("_", "-"))
@@ -88,6 +99,166 @@ class PipesK8sPodLogsMessageReader(PipesMessageReader):
             for log_line in log_chunk.split("\n"):
                 extract_message_or_forward_to_stdout(handler, log_line)
 
+    @contextmanager
+    def async_consume_pod_logs(
+        self,
+        context: Union[OpExecutionContext, AssetExecutionContext],
+        core_api: kubernetes.client.CoreV1Api,
+        pod_name: str,
+        namespace: str,
+    ) -> Generator:
+        """Consume all logs from all containers within the pod.
+
+        Args:
+            context (Union[OpExecutionContext, AssetExecutionContext]): The execution context.
+            core_api: The k8s core API.
+            pod_name: The pod to collect logs from.
+            namespace: The namespace to collect logs from.
+
+        """
+        handler = check.not_none(
+            self._handler, "can only consume logs within scope of context manager"
+        )
+        pods = core_api.list_namespaced_pod(
+            namespace=namespace, field_selector=f"metadata.name={pod_name}"
+        ).items
+
+        containers = []
+        # https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.18/#containerstatus-v1-core
+        for pod in pods:
+            if pod.status.init_container_statuses:
+                containers.extend(
+                    [
+                        container_status.name
+                        for container_status in pod.status.init_container_statuses
+                    ]
+                )
+
+            if pod.status.container_statuses:
+                containers.extend(
+                    [container_status.name for container_status in pod.status.container_statuses]
+                )
+
+        pod_exit_event = threading.Event()
+
+        logger = context.log.getChild("consume_pod_logs")
+        logger.setLevel(logging.WARNING)
+
+        with merge_streams(
+            streams={
+                f"{pod_name}:{container}": self._extract_logs(
+                    pod_exit_event=pod_exit_event,
+                    read_namespaced_pod_log=core_api.read_namespaced_pod_log,
+                    list_namespaced_pod=core_api.list_namespaced_pod,
+                    pod_name=pod_name,
+                    namespace=namespace,
+                    container=container,
+                    logger=logger.getChild(f"_extract_logs({container})"),
+                )
+                for container in containers
+            },
+            log_handler=lambda log_line: extract_message_or_forward_to_stdout(handler, log_line),
+            stream_processor=_process_log_stream,
+            logger=logger,
+        ):
+            yield
+            logger.info("Setting the pod exit event to do the cleanup of the streams")
+            pod_exit_event.set()
+
+    def _extract_logs(
+        self,
+        pod_exit_event: threading.Event,
+        read_namespaced_pod_log: Callable,
+        list_namespaced_pod: Callable,
+        pod_name: str,
+        namespace: str,
+        container: str,
+        logger: logging.Logger,
+        max_attempts: int = 3,
+        sleep_between_attempts: float = 0.5,
+        sleeper: Callable = time.sleep,
+    ) -> Generator:
+        """Return the streams of the Kubernetes logs with the appropriate buffer time.
+
+        Args:
+            pod_exit_event (threading.Event): The threading event that indicates to the
+                log reading thread that the pod has exited
+            read_namespaced_pod_log (kubernetes.client.CoreV1Api): The Kubernetes CoreV1Api client function for reading
+                logs.
+            list_namespaced_pod (kubernetes.client.CoreV1Api): The Kubernetes CoreV1Api client function for listing
+                pods and their state.
+            pod_name (str): The name of the Pipes Pod
+            namespace (str): The namespace the pod lives in.
+            container (str): The container to read logs from.
+            logger (logging.Logger): A logger instance for diagnostic logs.
+            max_attempts (int): The number of attempts to read logs in the beginning in
+                case we get a failure due to pod still starting.
+            sleep_between_attempts (float): Sleep between attempts in the beginning.
+            sleeper (Callable): The time.sleep equivalent.
+
+        Yields:
+            The Kubernetes pod log stream generator
+
+        """
+        # Yield the actual stream here to hide implementation detail from caller
+        # If readiness/liveness probes aren't configured
+        # pods can reach the "Ready" state from the API perspective
+        # but still reject incoming communication
+        attempt = 0
+        common_args = {
+            "name": pod_name,
+            "namespace": namespace,
+            "container": container,
+            "_preload_content": False,  # avoid JSON processing
+            "timestamps": True,  # Include timestamps for ordering and deduplication
+            "follow": True,
+        }
+
+        # Attempt to get the stream for the first time
+        while attempt < max_attempts:
+            try:
+                yield read_namespaced_pod_log(since_seconds=3600, **common_args).stream()
+                break
+            except kubernetes.client.ApiException as e:
+                if e.status in ["400", 400] and "PodInitializing" in str(e):
+                    # PodInitializing cannot accept log consumption
+                    sleeper(sleep_between_attempts)
+                    sleep_between_attempts *= 2  # exponential backoff
+                    attempt += 1
+                    continue
+
+        # After stream is initially yielded in above loop this while loop is a safeguard against the
+        # stream ending while the pod has not exitted. If so, we need to refresh the stream.
+        while not pod_exit_event.is_set():
+            # List the pods now and then use the status to decide whether we should exit
+            pods = list_namespaced_pod(
+                namespace=namespace, field_selector=f"metadata.name={pod_name}"
+            ).items
+
+            try:
+                yield read_namespaced_pod_log(since_seconds=5, **common_args).stream()
+            except Exception:
+                logger.exception(f"{container}: exception in getting logs")
+                break
+
+            # The logs are still available once the pod has exited and the above call will succeed, we add this extra
+            # statement where we will exit if the status of the container was terminated before we read the logs. That
+            # ensures that we get all of the logs (the merge_streams will deduplicate them) and we don't waste CPU
+            # cycles whilst trying to get more logs.
+            pod = pods[0] if pods else None
+            if pod is None:
+                break
+
+            all_statuses = []
+            all_statuses.extend(pod.status.init_container_statuses or [])
+            all_statuses.extend(pod.status.container_statuses or [])
+            if not all_statuses:
+                break
+
+            state_by_name = {status.name: status.state for status in all_statuses}
+            if state_by_name[container].terminated is not None:
+                break
+
     def no_messages_debug_text(self) -> str:
         return "Attempted to read messages by extracting them from kubernetes pod logs directly."
 
@@ -191,7 +362,7 @@ class PipesK8sClient(PipesClient, TreatAsResourceParam):
     def run(
         self,
         *,
-        context: OpExecutionContext,
+        context: Union[OpExecutionContext, AssetExecutionContext],
         extras: Optional[PipesExtras] = None,
         image: Optional[str] = None,
         command: Optional[Union[str, Sequence[str]]] = None,
@@ -200,10 +371,13 @@ class PipesK8sClient(PipesClient, TreatAsResourceParam):
         base_pod_meta: Optional[Mapping[str, Any]] = None,
         base_pod_spec: Optional[Mapping[str, Any]] = None,
         ignore_containers: Optional[Set] = None,
+        enable_multi_container_logs: bool = False,
     ) -> PipesClientCompletedInvocation:
         """Publish a kubernetes pod and wait for it to complete, enriched with the pipes protocol.
 
         Args:
+            context (Union[OpExecutionContext, AssetExecutionContext]):
+                The execution context.
             image (Optional[str]):
                 The image to set the first container in the pod spec to use.
             command (Optional[Union[str, Sequence[str]]]):
@@ -214,11 +388,11 @@ class PipesK8sClient(PipesClient, TreatAsResourceParam):
             env (Optional[Mapping[str,str]]):
                 A mapping of environment variable names to values to set on the first
                 container in the pod spec, on top of those configured on resource.
-            base_pod_meta (Optional[Mapping[str, Any]]:
+            base_pod_meta (Optional[Mapping[str, Any]]):
                 Raw k8s config for the k8s pod's metadata
                 (https://kubernetes.io/docs/reference/kubernetes-api/common-definitions/object-meta/#ObjectMeta)
                 Keys can either snake_case or camelCase. The name value will be overridden.
-            base_pod_spec (Optional[Mapping[str, Any]]:
+            base_pod_spec (Optional[Mapping[str, Any]]):
                 Raw k8s config for the k8s pod's pod spec
                 (https://kubernetes.io/docs/reference/kubernetes-api/workload-resources/pod-v1/#PodSpec).
                 Keys can either snake_case or camelCase. The dagster context will be readable
@@ -232,10 +406,12 @@ class PipesK8sClient(PipesClient, TreatAsResourceParam):
                 Override the default ext protocol message reader.
             ignore_containers (Optional[Set]): Ignore certain containers from waiting for termination. Defaults to
                 None.
+            enable_multi_container_logs (bool): Whether or not to enable multi-container log consumption.
 
         Returns:
             PipesClientCompletedInvocation: Wrapper containing results reported by the external
                 process.
+
         """
         self._load_k8s_config()
         client = DagsterKubernetesClient.production_client()
@@ -262,31 +438,79 @@ class PipesK8sClient(PipesClient, TreatAsResourceParam):
             )
             client.core_api.create_namespaced_pod(namespace, pod_body)
             try:
-                # if were doing direct pod reading, wait for pod to start and then stream logs out
-                if isinstance(self.message_reader, PipesK8sPodLogsMessageReader):
+                # Consume pod logs if possible
+                with self.consume_pod_logs(
+                    context=context,
+                    client=client,
+                    namespace=namespace,
+                    pod_name=pod_name,
+                    enable_multi_container_logs=enable_multi_container_logs,
+                ):
+                    # We need to wait for the pod to start up so that the log streaming is successful afterwards.
                     client.wait_for_pod(
                         pod_name,
                         namespace,
-                        wait_for_state=WaitForPodState.Ready,
+                        wait_for_state=WaitForPodState.Terminated,
+                        ignore_containers=ignore_containers,
                         wait_time_between_attempts=self.poll_interval,
                     )
-                    self.message_reader.consume_pod_logs(
-                        core_api=client.core_api,
-                        pod_name=pod_name,
-                        namespace=namespace,
-                    )
-
-                client.wait_for_pod(
-                    pod_name,
-                    namespace,
-                    wait_for_state=WaitForPodState.Terminated,
-                    ignore_containers=ignore_containers,
-                    wait_time_between_attempts=self.poll_interval,
-                )
             finally:
                 client.core_api.delete_namespaced_pod(pod_name, namespace)
+
         return PipesClientCompletedInvocation(pipes_session)
 
+    @contextmanager
+    def consume_pod_logs(
+        self,
+        context: Union[OpExecutionContext, AssetExecutionContext],
+        client: DagsterKubernetesClient,
+        namespace: str,
+        pod_name: str,
+        enable_multi_container_logs: bool = False,
+    ) -> Iterator:
+        """Consume pod logs in the background if possible simple context manager to setup pod log consumption.
+
+        This will be a no-op if the message_reader is of the wrong type.
+
+        Args:
+            context (Union[OpExecutionContext, AssetExecutionContext]): The execution context.
+            client (kubernetes.client): _description_
+            namespace (str): The namespace the pod lives in
+            pod_name (str): The name of the Pipes Pod
+            enable_multi_container_logs (bool): Whether or not to enable multi-container log consumption
+
+        """
+        if isinstance(self.message_reader, PipesK8sPodLogsMessageReader):
+            # We need to wait for the pod to start up so that the log streaming is successful afterwards.
+            client.wait_for_pod(
+                pod_name,
+                namespace,
+                wait_for_state=WaitForPodState.Ready,
+                wait_time_between_attempts=self.poll_interval,
+                # After init container gains a status in the first while loop, there is still a check for
+                # the ready state in the second while loop, which respects the below timeout only.
+                # Very rarely, the pod will be Evicted there and we have to wait the default, unless set.
+                wait_timeout=WAIT_TIMEOUT_FOR_READY,
+            )
+
+            if enable_multi_container_logs:
+                with self.message_reader.async_consume_pod_logs(
+                    context=context,
+                    core_api=client.core_api,
+                    namespace=namespace,
+                    pod_name=pod_name,
+                ):
+                    yield
+                    return
+            else:
+                self.message_reader.consume_pod_logs(
+                    core_api=client.core_api,
+                    namespace=namespace,
+                    pod_name=pod_name,
+                )
+
+        yield
+
 
 def _detect_current_namespace(
     kubeconfig_file: Optional[str], namespace_secret_path: Path = _NAMESPACE_SECRET_PATH
@@ -423,3 +647,84 @@ def build_pod_body(
             "spec": spec,
         },
     )
+
+
+def _process_log_stream(stream: Iterator[bytes]) -> Iterator[LogItem]:
+    """This expects the logs to be of the format b'<timestamp> <msg>' and only the
+    '<msg>' is forwarded to Dagster. If the <timestamp> is not there then the lines
+    will be joined together. There is a limitation that the first item in the stream
+    needs to always contain a timestamp as the first element.
+
+    The timestamp is expected to be in '2024-03-22T02:17:29.885548Z' format and
+    if the subsecond part will be truncated to microseconds.
+
+    If we fail parsing the timestamp, then the priority will be set to zero in
+    order to not drop any log items.
+
+    Args:
+        stream (Iterator[bytes]): A stream of log chunks
+
+    Yields:
+        Iterator[LogItem]: A log containing the timestamp and msg
+    """
+    timestamp = ""
+    log = ""
+
+    for log_chunk in stream:
+        for line in log_chunk.decode("utf-8").split("\n"):
+            maybe_timestamp, _, tail = line.partition(" ")
+            if not timestamp:
+                # The first item in the stream will always have a timestamp.
+                timestamp = maybe_timestamp
+                log = tail
+            elif maybe_timestamp == timestamp:
+                # We have multiple messages with the same timestamp in this chunk, add them separated
+                # with a new line
+                log += f"\n{tail}"
+            elif not (
+                len(maybe_timestamp) == len(timestamp) and _is_kube_timestamp(maybe_timestamp)
+            ):
+                # The line is continuation of a long line that got truncated and thus doesn't
+                # have a timestamp in the beginning of the line.
+                # Since all timestamps in the RFC format returned by Kubernetes have the same
+                # length (when represented as strings) we know that the value won't be a timestamp
+                # if the string lengths differ, however if they do not differ, we need to parse the
+                # timestamp.
+                log += line
+            else:
+                # New log line has been observed, send in the next cycle
+                yield LogItem(timestamp=timestamp, log=log)
+                timestamp = maybe_timestamp
+                log = tail
+
+    # Send the last message that we were building
+    if log or timestamp:
+        yield LogItem(timestamp=timestamp, log=log)
+
+
+def _is_kube_timestamp(maybe_timestamp: str) -> bool:
+    # This extra stripping logic is necessary, as Python's strptime fn doesn't
+    # handle valid ISO 8601 timestamps with nanoseconds which we receive in k8s
+    # e.g. 2024-03-22T02:17:29.185548486Z
+
+    # This is likely fine. We're just trying to confirm whether or not it's a
+    # valid timestamp, not trying to parse it with full correctness.
+    if maybe_timestamp.endswith("Z"):
+        maybe_timestamp = maybe_timestamp[:-1]  # Strip the "Z"
+        if "." in maybe_timestamp:
+            # Split at the decimal point to isolate the fractional seconds
+            date_part, frac_part = maybe_timestamp.split(".")
+            maybe_timestamp = f"{date_part}.{frac_part[:6]}Z"
+        else:
+            maybe_timestamp = f"{maybe_timestamp}Z"  # Add the "Z" back if no fractional part
+        try:
+            datetime.strptime(maybe_timestamp, "%Y-%m-%dT%H:%M:%S.%fZ")
+            return True
+        except ValueError:
+            return False
+    else:
+        try:
+            datetime.strptime(maybe_timestamp, "%Y-%m-%dT%H:%M:%S%z")
+            return True
+        except ValueError:
+            return False

dagster_k8s/version.py

@@ -1 +1 @@
-__version__ = "0.24.12"
+__version__ = "0.25.0"

{dagster_k8s-0.24.12.dist-info → dagster_k8s-0.25.0.dist-info}/METADATA

@@ -1,21 +1,20 @@
 Metadata-Version: 2.1
 Name: dagster-k8s
-Version: 0.24.12
+Version: 0.25.0
 Summary: A Dagster integration for k8s
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-k8s
 Author: Dagster Labs
 Author-email: hello@dagsterlabs.com
 License: Apache-2.0
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.8,<3.13
+Requires-Python: >=3.9,<3.13
 License-File: LICENSE
-Requires-Dist: dagster ==1.8.12
+Requires-Dist: dagster ==1.9.0
 Requires-Dist: kubernetes <32
 Requires-Dist: google-auth !=2.23.1
 

{dagster_k8s-0.24.12.dist-info → dagster_k8s-0.25.0.dist-info}/RECORD

@@ -6,15 +6,15 @@ dagster_k8s/job.py,sha256=-GzUcz80cwF9NwDAv3QdE6GezSSpUzF0kLRkj1oTO34,43046
 dagster_k8s/kubernetes_version.py,sha256=jIBF12yvVweYUCmhC5AJ2Lb1JPcHxJfohJ1_hHfzS4o,38
 dagster_k8s/launcher.py,sha256=Mfs4lD8pSjqC0S39d0csPJk6g-ieVuiyv680nWtSezw,16621
 dagster_k8s/models.py,sha256=OPksL8WjsdqJtGcw4WLkvsnrm1WA6vCUjqHL9dA2oYQ,6270
-dagster_k8s/pipes.py,sha256=YRPvEACWI0lmkodNnUNF26knJghQGPmTAcbY7rTRPY8,17247
+dagster_k8s/pipes.py,sha256=kRHLYJfZi23GJe1_TW3Z3FglEpMrMLT43UV7YpTmKr8,30439
 dagster_k8s/py.typed,sha256=la67KBlbjXN-_-DfGNcdOcjYumVpKG_Tkw-8n5dnGB4,8
 dagster_k8s/test.py,sha256=cNtcbzxytiZtd01wY5ip7KPi01y0BUQuQhohoIfAFUM,684
 dagster_k8s/utils.py,sha256=c1bHqh5f1p5RZ0JCT6WEbPPjDvbgUp3pl4nYZRaaI4s,786
-dagster_k8s/version.py,sha256=8e_5IttRuoue-uppwtfAsOgcCm05WlECsinJs_aHSOc,24
+dagster_k8s/version.py,sha256=Mu4JbSLl5nr-J2figk5hmW2mrw4skf_oeIzxbnpcgwY,23
 dagster_k8s/ops/__init__.py,sha256=ur-9GrE_DRfnsFCpYan03qOY9cWbjagC8KHZFZuiCmc,113
 dagster_k8s/ops/k8s_job_op.py,sha256=uLpQE3ZJbXDiGyyRZydyp-oDQn5XvQuhRcUwd2qiJho,21124
-dagster_k8s-0.24.12.dist-info/LICENSE,sha256=TMatHW4_G9ldRdodEAp-l2Xa2WvsdeOh60E3v1R2jis,11349
-dagster_k8s-0.24.12.dist-info/METADATA,sha256=gUddbr3Tv7LkSU68woFy7J5sdWIKfh_THgi9nC5mKAM,784
-dagster_k8s-0.24.12.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-dagster_k8s-0.24.12.dist-info/top_level.txt,sha256=wFPjskoWPlk2hOLugYCaoZhSiZdUcbCA1QZe9I4dals,12
-dagster_k8s-0.24.12.dist-info/RECORD,,
+dagster_k8s-0.25.0.dist-info/LICENSE,sha256=TMatHW4_G9ldRdodEAp-l2Xa2WvsdeOh60E3v1R2jis,11349
+dagster_k8s-0.25.0.dist-info/METADATA,sha256=YEeTIilD0QIF7_6WZKxleeR6-cmuI1KHXecZIcJ8ics,732
+dagster_k8s-0.25.0.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+dagster_k8s-0.25.0.dist-info/top_level.txt,sha256=wFPjskoWPlk2hOLugYCaoZhSiZdUcbCA1QZe9I4dals,12
+dagster_k8s-0.25.0.dist-info/RECORD,,