skypilot-nightly 1.0.0.dev20251203__py3-none-any.whl → 1.0.0.dev20260112__py3-none-any.whl

sky/utils/kubernetes/ssh-tunnel.sh

@@ -1,379 +1,10 @@
 #!/bin/bash
-# ssh-tunnel.sh - SSH tunnel script for Kubernetes API access
-# Used as kubectl exec credential plugin to establish SSH tunnel on demand.
-# Returns a valid credential format for kubectl with expiration. The expiration
-# is calculated based on the TTL argument and is required to force kubectl to
-# check the tunnel status frequently.
 
-# Usage: ssh-tunnel.sh --host HOST [--user USER] [--use-ssh-config] [--ssh-key KEY] [--context CONTEXT] [--port PORT] [--ttl SECONDS]
+# This redirect stub is needed because we use this script in the
+# exec auth section when creating our kubeconfig. Therefore, node pools
+# launched in older versions of SkyPilot will have kubeconfigs pointing
+# to this path.
 
-# Default time-to-live for credential in seconds
-# This forces kubectl to check the tunnel status frequently
-TTL_SECONDS=30
-
-# Parse arguments
-USE_SSH_CONFIG=0
-SSH_KEY=""
-CONTEXT=""
-HOST=""
-USER=""
-PORT=6443  # Default port if not specified
-
-# Debug log to ~/.sky/ssh_node_pools_info/$CONTEXT-tunnel.log
-debug_log() {
-    local message="$(date): $1"
-    echo "$message" >> "$LOG_FILE"
-}
-
-# Generate expiration timestamp for credential
-generate_expiration_timestamp() {
-  # Try macOS date format first, fallback to Linux format
-  date -u -v+${TTL_SECONDS}S +"%Y-%m-%dT%H:%M:%SZ" 2>/dev/null || date -u -d "+${TTL_SECONDS} seconds" +"%Y-%m-%dT%H:%M:%SZ"
-}
-
-# Acquire the lock, return 0 if successful, 1 if another process is already holding the lock
-acquire_lock() {
-  # Check for flock command
-  if ! command -v flock >/dev/null 2>&1; then
-    debug_log "flock command not available, using alternative lock mechanism"
-    # Simple file-based locking
-    if [ -f "$LOCK_FILE" ]; then
-      lock_pid=$(cat "$LOCK_FILE" 2>/dev/null)
-      if [ -n "$lock_pid" ] && kill -0 "$lock_pid" 2>/dev/null; then
-        debug_log "Another process ($lock_pid) is starting the tunnel, waiting briefly"
-        return 1
-      else
-        # Stale lock file
-        debug_log "Removing stale lock file"
-        rm -f "$LOCK_FILE"
-      fi
-    fi
-    # Create our lock
-    echo $$ > "$LOCK_FILE"
-    return 0
-  else
-    # Use flock for better locking
-    exec 9>"$LOCK_FILE"
-    if ! flock -n 9; then
-      debug_log "Another process is starting the tunnel, waiting briefly"
-      return 1
-    fi
-    return 0
-  fi
-}
-
-# Release the lock
-release_lock() {
-  if command -v flock >/dev/null 2>&1; then
-    # Using flock
-    exec 9>&- # Close file descriptor to release lock
-  else
-    # Using simple lock
-    rm -f "$LOCK_FILE"
-  fi
-  debug_log "Lock released"
-}
-
-# Generate SSH command based on available tools and parameters
-generate_ssh_command() {
-  # Check for autossh
-  if ! command -v autossh >/dev/null 2>&1; then
-    debug_log "WARNING: autossh is not installed but recommended for reliable SSH tunnels"
-    debug_log "Install autossh: brew install autossh (macOS), apt-get install autossh (Ubuntu/Debian)"
-    
-    # Fall back to regular ssh
-    if [[ $USE_SSH_CONFIG -eq 1 ]]; then
-      SSH_CMD=("ssh" "-o" "ServerAliveInterval=30" "-o" "ServerAliveCountMax=3" "-o" "ExitOnForwardFailure=yes" "-L" "$PORT:127.0.0.1:6443" "-N" "$HOST")
-    else
-      SSH_CMD=("ssh" "-o" "StrictHostKeyChecking=no" "-o" "IdentitiesOnly=yes" "-o" "ServerAliveInterval=30" "-o" "ServerAliveCountMax=3" "-o" "ExitOnForwardFailure=yes" "-L" "$PORT:127.0.0.1:6443" "-N")
-      
-      # Add SSH key if provided
-      if [[ -n "$SSH_KEY" ]]; then
-        SSH_CMD+=("-i" "$SSH_KEY")
-      fi
-      
-      # Add user@host
-      SSH_CMD+=("$USER@$HOST")
-    fi
-  else
-    # Configure autossh
-    if [[ $USE_SSH_CONFIG -eq 1 ]]; then
-      SSH_CMD=("autossh" "-M" "0" "-o" "ServerAliveInterval=30" "-o" "ServerAliveCountMax=3" "-o" "ExitOnForwardFailure=yes" "-L" "$PORT:127.0.0.1:6443" "-N" "$HOST")
-    else
-      SSH_CMD=("autossh" "-M" "0" "-o" "StrictHostKeyChecking=no" "-o" "IdentitiesOnly=yes" "-o" "ServerAliveInterval=30" "-o" "ServerAliveCountMax=3" "-o" "ExitOnForwardFailure=yes" "-L" "$PORT:127.0.0.1:6443" "-N")
-      
-      # Add SSH key if provided
-      if [[ -n "$SSH_KEY" ]]; then
-        SSH_CMD+=("-i" "$SSH_KEY")
-      fi
-      
-      # Add user@host
-      SSH_CMD+=("$USER@$HOST")
-    fi
-  fi
-}
-
-# Function to read certificate files if they exist
-read_certificate_data() {
-  local client_cert_file="$TUNNEL_DIR/$CONTEXT-cert.pem"
-  local client_key_file="$TUNNEL_DIR/$CONTEXT-key.pem"
-  local cert_data=""
-  local key_data=""
-  
-  if [[ -f "$client_cert_file" ]]; then
-    # Read the certificate file as is - it's already in PEM format
-    cert_data=$(cat "$client_cert_file")
-    debug_log "Found client certificate data for context $CONTEXT"
-    
-    # Log the first and last few characters to verify PEM format
-    local cert_start=$(head -1 "$client_cert_file")
-    local cert_end=$(tail -1 "$client_cert_file")
-    debug_log "Certificate starts with: $cert_start"
-    debug_log "Certificate ends with: $cert_end"
-    
-    # Check if it has proper PEM format
-    if ! grep -q "BEGIN CERTIFICATE" "$client_cert_file" || ! grep -q "END CERTIFICATE" "$client_cert_file"; then
-      debug_log "WARNING: Certificate file may not be in proper PEM format"
-      # Try to fix it if needed
-      if ! grep -q "BEGIN CERTIFICATE" "$client_cert_file"; then
-        echo "-----BEGIN CERTIFICATE-----" > "$client_cert_file.fixed"
-        cat "$client_cert_file" >> "$client_cert_file.fixed"
-        echo "-----END CERTIFICATE-----" >> "$client_cert_file.fixed"
-        mv "$client_cert_file.fixed" "$client_cert_file"
-        cert_data=$(cat "$client_cert_file")
-        debug_log "Fixed certificate format by adding BEGIN/END markers"
-      fi
-    fi
-  fi
-  
-  if [[ -f "$client_key_file" ]]; then
-    # Read the key file as is - it's already in PEM format
-    key_data=$(cat "$client_key_file")
-    debug_log "Found client key data for context $CONTEXT"
-    
-    # Log the first and last few characters to verify PEM format
-    local key_start=$(head -1 "$client_key_file")
-    local key_end=$(tail -1 "$client_key_file")
-    debug_log "Key starts with: $key_start"
-    debug_log "Key ends with: $key_end"
-    
-    # Check if it has proper PEM format
-    if ! grep -q "BEGIN" "$client_key_file" || ! grep -q "END" "$client_key_file"; then
-      debug_log "WARNING: Key file may not be in proper PEM format"
-      # Try to fix it if needed
-      if ! grep -q "BEGIN" "$client_key_file"; then
-        echo "-----BEGIN PRIVATE KEY-----" > "$client_key_file.fixed"
-        cat "$client_key_file" >> "$client_key_file.fixed"
-        echo "-----END PRIVATE KEY-----" >> "$client_key_file.fixed"
-        mv "$client_key_file.fixed" "$client_key_file"
-        key_data=$(cat "$client_key_file")
-        debug_log "Fixed key format by adding BEGIN/END markers"
-      fi
-    fi
-  fi
-  
-  echo "$cert_data:$key_data"
-}
-
-# Function to generate credentials JSON
-generate_credentials_json() {
-  local expiration_time=$(generate_expiration_timestamp)
-  local cert_bundle=$(read_certificate_data)
-  local client_cert_data=${cert_bundle%:*}
-  local client_key_data=${cert_bundle#*:}
-  
-  if [[ -n "$client_cert_data" && -n "$client_key_data" ]]; then
-    # Debug the certificate data
-    debug_log "Certificate data length: $(echo -n "$client_cert_data" | wc -c) bytes"
-    debug_log "Key data length: $(echo -n "$client_key_data" | wc -c) bytes"
-    
-    # Check if we can create proper JSON with `jq`
-    if ! command -v jq &>/dev/null; then
-      echo "jq is not installed. Please install jq to use this script." >&2
-      exit 1
-    fi
-    debug_log "Using jq for JSON formatting"
-    
-    # Create a temporary file for the JSON output to avoid shell escaping issues
-    local TEMP_JSON_FILE=$(mktemp)
-    
-    # Write the JSON to the temporary file using jq for proper JSON formatting
-    cat > "$TEMP_JSON_FILE" << EOL
-{
-  "apiVersion": "client.authentication.k8s.io/v1beta1",
-  "kind": "ExecCredential",
-  "status": {
-    "clientCertificateData": $(printf '%s' "$client_cert_data" | jq -R -s .),
-    "clientKeyData": $(printf '%s' "$client_key_data" | jq -R -s .),
-    "expirationTimestamp": "$expiration_time"
-  }
-}
-EOL
-      
-    # Read the JSON from the file
-    local json_response=$(cat "$TEMP_JSON_FILE")
-    
-    # Clean up
-    rm -f "$TEMP_JSON_FILE"
-    
-    # Output the JSON
-    echo "$json_response"
-  else
-    # Fallback to token-based credential for tunnel-only authentication
-    echo "{\"apiVersion\":\"client.authentication.k8s.io/v1beta1\",\"kind\":\"ExecCredential\",\"status\":{\"token\":\"k8s-ssh-tunnel-token\",\"expirationTimestamp\":\"$expiration_time\"}}"
-  fi
-}
-
-while [[ $# -gt 0 ]]; do
-  case $1 in
-    --use-ssh-config)
-      USE_SSH_CONFIG=1
-      shift
-      ;;
-    --ssh-key)
-      SSH_KEY="$2"
-      shift 2
-      ;;
-    --context)
-      CONTEXT="$2"
-      shift 2
-      ;;
-    --port)
-      PORT="$2"
-      shift 2
-      ;;
-    --host)
-      HOST="$2"
-      shift 2
-      ;;
-    --user)
-      USER="$2"
-      shift 2
-      ;;
-    --ttl)
-      TTL_SECONDS="$2"
-      shift 2
-      ;;
-    *)
-      echo "Unknown parameter: $1" >&2
-      exit 1
-      ;;
-  esac
-done
-
-# Validate required parameters
-if [[ -z "$HOST" ]]; then
-  echo "Error: --host parameter is required" >&2
-  exit 1
-fi
-
-# Setup directories
-TUNNEL_DIR="$HOME/.sky/ssh_node_pools_info"
-mkdir -p "$TUNNEL_DIR"
-
-# Get context name for PID file
-if [[ -z "$CONTEXT" ]]; then
-  CONTEXT="default"
-fi
-
-PID_FILE="$TUNNEL_DIR/$CONTEXT-tunnel.pid"
-LOG_FILE="$TUNNEL_DIR/$CONTEXT-tunnel.log"
-LOCK_FILE="$TUNNEL_DIR/$CONTEXT-tunnel.lock"
-
-debug_log "Starting ssh-tunnel.sh for context $CONTEXT, host $HOST, port $PORT"
-debug_log "SSH Config: $USE_SSH_CONFIG, User: $USER, TTL: ${TTL_SECONDS}s"
-
-# Check if specified port is already in use (tunnel may be running)
-if nc -z 127.0.0.1 "$PORT" 2>/dev/null; then
-  debug_log "Port $PORT already in use, checking if it's our tunnel"
-  
-  # Check if there's a PID file and if that process is running
-  if [[ -f "$PID_FILE" ]]; then
-    OLD_PID=$(cat "$PID_FILE")
-    if kill -0 "$OLD_PID" 2>/dev/null; then
-      debug_log "Tunnel appears to be running with PID $OLD_PID"
-    else
-      debug_log "PID file exists but process $OLD_PID is not running"
-    fi
-  else
-    debug_log "Port $PORT is in use but no PID file exists"
-  fi
-  
-  # Return valid credential format for kubectl with expiration
-  generate_credentials_json
-  exit 0
-fi
-
-# Try to acquire the lock
-if ! acquire_lock; then
-  # Wait briefly for the tunnel to be established
-  for i in {1..10}; do
-    if nc -z 127.0.0.1 "$PORT" 2>/dev/null; then
-      debug_log "Tunnel is now active"
-      
-      # Return valid credential format for kubectl with expiration
-      generate_credentials_json
-      exit 0
-    fi
-    sleep 0.2
-  done
-  debug_log "Waited for tunnel but port $PORT still not available"
-fi
-
-# Check if we have a PID file with running process
-if [[ -f "$PID_FILE" ]]; then
-  OLD_PID=$(cat "$PID_FILE")
-  if kill -0 "$OLD_PID" 2>/dev/null; then
-    # Process exists but port isn't open - something's wrong, kill it
-    kill "$OLD_PID" 2>/dev/null
-    debug_log "Killed stale tunnel process $OLD_PID"
-  else
-    debug_log "PID file exists but process $OLD_PID is not running anymore"
-  fi
-  # Remove the stale PID file
-  rm -f "$PID_FILE"
-fi
-
-# Generate the SSH command
-generate_ssh_command
-
-debug_log "Starting SSH tunnel: ${SSH_CMD[*]}"
-
-# Start the tunnel in foreground and wait for it to establish
-"${SSH_CMD[@]}" >> "$LOG_FILE" 2>&1 &
-TUNNEL_PID=$!
-
-# Save PID
-echo $TUNNEL_PID > "$PID_FILE"
-debug_log "Tunnel started with PID $TUNNEL_PID"
-
-# Wait for tunnel to establish
-tunnel_up=0
-for i in {1..20}; do
-  if nc -z 127.0.0.1 "$PORT" 2>/dev/null; then
-    debug_log "Tunnel established successfully on port $PORT"
-    tunnel_up=1
-    break
-  fi
-  sleep 0.2
-done
-
-# Clean up lock file
-release_lock
-
-# Check if the tunnel process is still running
-if ! kill -0 $TUNNEL_PID 2>/dev/null; then
-  debug_log "ERROR: Tunnel process exited unexpectedly! Check logs for details"
-  if [[ -f "$PID_FILE" ]]; then
-    rm -f "$PID_FILE"
-  fi
-  # Return error in case of tunnel failure
-  echo "Failed to establish SSH tunnel. See $TUNNEL_DIR/$CONTEXT-tunnel.log for details." >&2
-  exit 1
-elif [[ $tunnel_up -eq 0 ]]; then
-  debug_log "WARNING: Tunnel process is running but port $PORT is not responding"
-fi
-
-# Return valid credential format with certificates if available
-generate_credentials_json
-exit 0 
+# TODO (kyuds): remove this script after v0.13.0. Kept here for backwards compat.
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+exec "$SCRIPT_DIR/../../ssh_node_pools/deploy/tunnel/ssh-tunnel.sh" "$@"

sky/utils/plugin_extensions/__init__.py

@@ -0,0 +1,14 @@
+"""Plugin extensions module.
+
+This module provides extension points that plugins can hook into to provide
+custom functionality.
+"""
+from sky.utils.plugin_extensions.external_failure_source import (
+    ExternalClusterFailure)
+from sky.utils.plugin_extensions.external_failure_source import (
+    ExternalFailureSource)
+
+__all__ = [
+    'ExternalClusterFailure',
+    'ExternalFailureSource',
+]

sky/utils/plugin_extensions/external_failure_source.py

@@ -0,0 +1,176 @@
+"""External failure source interface for plugins.
+
+This module provides an extension point that allows plugins to provide
+cluster failure tracking functionality. By default, no-op implementations
+are used. Plugins can register their own implementations to provide actual
+failure tracking.
+
+Example usage in a plugin:
+    from sky.utils.plugin_extensions import ExternalFailureSource
+
+    # Register custom failure source
+    ExternalFailureSource.register(
+        get_failures=my_get_cluster_failures,
+        clear_failures=my_clear_cluster_failures,
+    )
+
+Example usage in core SkyPilot:
+    from sky.utils.plugin_extensions import ExternalFailureSource
+
+    # Get failures for a cluster
+    failures = ExternalFailureSource.get(cluster_hash='abc123')
+
+    # Clear failures for a cluster
+    cleared = ExternalFailureSource.clear(cluster_name='my-cluster')
+"""
+import dataclasses
+from typing import Any, Dict, List, Optional, Protocol
+
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
+
+
+@dataclasses.dataclass
+class ExternalClusterFailure:
+    """Represents a single cluster failure from an external source.
+
+    Attributes:
+        code: Machine-readable failure code (e.g. 'GPU_HARDWARE_FAILURE_XID_79')
+        reason: Human-readable description of the failure.
+    """
+    code: str
+    reason: str
+
+    @classmethod
+    def from_failure_list(
+            cls, failures: List[Dict[str,
+                                     Any]]) -> List['ExternalClusterFailure']:
+        """Create a list of ExternalClusterFailure from failure dicts.
+
+        Args:
+            failures: List of dicts with 'failure_mode' and 'failure_reason'
+                keys (as returned by ExternalFailureSource.get()).
+
+        Returns:
+            List of ExternalClusterFailure objects, one per failure.
+        """
+        return [
+            cls(code=f['failure_mode'], reason=f['failure_reason'])
+            for f in failures
+        ]
+
+
+# Protocol definitions for the failure source functions
+class GetClusterFailuresFunc(Protocol):
+    """Protocol for get_cluster_failures function."""
+
+    def __call__(self,
+                 cluster_hash: Optional[str] = None,
+                 cluster_name: Optional[str] = None) -> List[Dict[str, Any]]:
+        ...
+
+
+class ClearClusterFailuresFunc(Protocol):
+    """Protocol for clear_cluster_failures function."""
+
+    def __call__(self,
+                 cluster_hash: Optional[str] = None,
+                 cluster_name: Optional[str] = None) -> List[Dict[str, Any]]:
+        ...
+
+
+class ExternalFailureSource:
+    """Singleton class for external cluster failure source.
+
+    This class provides an extension point for plugins to register their own
+    cluster failure tracking implementations. By default, no-op implementations
+    are used that return empty lists.
+
+    Plugins can register their implementations during their install() phase,
+    and core SkyPilot code can use the get() and clear() methods to interact
+    with cluster failures without knowing which plugin (if any) is providing
+    the implementation.
+    """
+
+    _get_func: Optional[GetClusterFailuresFunc] = None
+    _clear_func: Optional[ClearClusterFailuresFunc] = None
+
+    @classmethod
+    def register(cls, get_failures: GetClusterFailuresFunc,
+                 clear_failures: ClearClusterFailuresFunc) -> None:
+        """Register an external failure source implementation.
+
+        This allows plugins to provide their own cluster failure tracking.
+        Only one external failure source can be registered at a time.
+
+        Args:
+            get_failures: Function to get active cluster failures.
+                Signature: (cluster_hash: Optional[str],
+                            cluster_name: Optional[str])
+                           -> List[Dict[str, Any]]
+                Returns list of dicts with keys: cluster_hash, failure_mode,
+                failure_reason, cleared_at.
+            clear_failures: Function to clear cluster failures.
+                Signature: (cluster_hash: Optional[str],
+                            cluster_name: Optional[str])
+                           -> List[Dict[str, Any]]
+                Returns list of dicts of the failures that were cleared.
+        """
+        cls._get_func = get_failures
+        cls._clear_func = clear_failures
+        logger.info('Registered external failure source')
+
+    @classmethod
+    def is_registered(cls) -> bool:
+        """Check if an external failure source is registered."""
+        return cls._get_func is not None and cls._clear_func is not None
+
+    @classmethod
+    def get(cls,
+            cluster_hash: Optional[str] = None,
+            cluster_name: Optional[str] = None) -> List[Dict[str, Any]]:
+        """Get active cluster failures from the registered failure source.
+
+        Args:
+            cluster_hash: Hash of the cluster to query failures for.
+            cluster_name: Name of the cluster to query failures for.
+
+        Returns:
+            List of dictionaries containing failure records.
+            Each dict contains: cluster_hash, failure_mode, failure_reason,
+            cleared_at. Returns empty list if no failure source is registered.
+        """
+        if cls._get_func is None:
+            return []
+        try:
+            # pylint: disable=not-callable
+            return cls._get_func(cluster_name=cluster_name,
+                                 cluster_hash=cluster_hash)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning(f'Failed to get cluster failures: {e}')
+            return []
+
+    @classmethod
+    def clear(cls,
+              cluster_hash: Optional[str] = None,
+              cluster_name: Optional[str] = None) -> List[Dict[str, Any]]:
+        """Clear cluster failures via the registered failure source.
+
+        Args:
+            cluster_hash: Hash of the cluster to clear failures for.
+            cluster_name: Name of the cluster to clear failures for.
+
+        Returns:
+            List of dictionaries containing the failure records that were
+            cleared. Returns empty list if no failure source is registered.
+        """
+        if cls._clear_func is None:
+            return []
+        try:
+            # pylint: disable=not-callable
+            return cls._clear_func(cluster_name=cluster_name,
+                                   cluster_hash=cluster_hash)
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning(f'Failed to clear cluster failures: {e}')
+            return []

sky/utils/resources_utils.py

@@ -183,7 +183,8 @@ def simplify_ports(ports: List[str]) -> List[str]:
 def format_resource(resource: 'resources_lib.Resources',
                     simplified_only: bool = False) -> Tuple[str, Optional[str]]:
     resource = resource.assert_launchable()
-    is_k8s = str(resource.cloud).lower() == 'kubernetes'
+    is_k8s = resource.cloud.canonical_name() == 'kubernetes'
+    vcpu, mem = None, None
     if resource.accelerators is None or is_k8s or not simplified_only:
         vcpu, mem = resource.cloud.get_vcpus_mem_from_instance_type(
             resource.instance_type)
@@ -198,18 +199,19 @@ def format_resource(resource: 'resources_lib.Resources',
 
     if (resource.accelerators is None or is_k8s):
         if vcpu is not None:
-            elements_simple.append(f'cpus={int(vcpu)}')
-            elements_full.append(f'cpus={int(vcpu)}')
+            elements_simple.append(f'cpus={common_utils.format_float(vcpu)}')
+            elements_full.append(f'cpus={common_utils.format_float(vcpu)}')
         if mem is not None:
-            elements_simple.append(f'mem={int(mem)}')
-            elements_full.append(f'mem={int(mem)}')
+            elements_simple.append(f'mem={common_utils.format_float(mem)}')
+            elements_full.append(f'mem={common_utils.format_float(mem)}')
     elif not simplified_only:
         if vcpu is not None:
-            elements_full.append(f'cpus={int(vcpu)}')
+            elements_full.append(f'cpus={common_utils.format_float(vcpu)}')
         if mem is not None:
-            elements_full.append(f'mem={int(mem)}')
+            elements_full.append(f'mem={common_utils.format_float(mem)}')
 
-    if not is_k8s:
+    is_slurm = resource.cloud.canonical_name() == 'slurm'
+    if not is_k8s and not is_slurm:
         instance_type_full = resource.instance_type
         instance_type_simple = common_utils.truncate_long_string(
             instance_type_full, 15)

sky/utils/rich_utils.py

@@ -362,14 +362,14 @@ def decode_rich_status(
                     # Replace `\r\n` with `\n`, as printing a line ends with
                     # `\r\n` in linux will cause the line to be empty.
                     line = line[:-2] + '\n'
-                is_payload, line = message_utils.decode_payload(
+                is_payload, decoded_line = message_utils.decode_payload(
                     line, raise_for_mismatch=False)
-                control = None
-                if is_payload:
-                    control, encoded_status = Control.decode(line)
-                if control is None:
+                if not is_payload:
                     yield line
                     continue
+                control, encoded_status = Control.decode(decoded_line)
+                if control is None:
+                    continue
 
                 if control == Control.RETRY:
                     raise exceptions.RequestInterruptedError(
@@ -481,15 +481,13 @@ async def decode_rich_status_async(
                     # Replace `\r\n` with `\n`, as printing a line ends with
                     # `\r\n` in linux will cause the line to be empty.
                     line = line[:-2] + '\n'
-                is_payload, line = message_utils.decode_payload(
+                is_payload, decoded_line = message_utils.decode_payload(
                     line, raise_for_mismatch=False)
-                if line is None:
+                if not is_payload:
+                    yield line
                     continue
-                control = None
-                if is_payload:
-                    control, encoded_status = Control.decode(line)
+                control, encoded_status = Control.decode(decoded_line)
                 if control is None:
-                    yield line
                     continue
 
                 if control == Control.RETRY: