konduktor-nightly 0.1.0.dev20250722105323__tar.gz → 0.1.0.dev20250723105251__tar.gz

{konduktor_nightly-0.1.0.dev20250722105323 → konduktor_nightly-0.1.0.dev20250723105251}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: konduktor-nightly
-Version: 0.1.0.dev20250722105323
+Version: 0.1.0.dev20250723105251
 Summary: GPU Cluster Health Management
 Author: Andrew Aikawa
 Author-email: asai@berkeley.edu

{konduktor_nightly-0.1.0.dev20250722105323 → konduktor_nightly-0.1.0.dev20250723105251}/konduktor/__init__.py

@@ -14,7 +14,7 @@ __all__ = [
 ]
 
 # Replaced with the current commit when building the wheels.
-_KONDUKTOR_COMMIT_SHA = '84818710a16e0a0515fbbd7878395fca37cf94f7'
+_KONDUKTOR_COMMIT_SHA = '4572b2fb2c49fb54573c080a71ab9d982ced5bf2'
 os.makedirs(os.path.expanduser('~/.konduktor'), exist_ok=True)
 
 
@@ -48,5 +48,5 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev0.1.0.dev20250722105323'
+__version__ = '1.0.0.dev0.1.0.dev20250723105251'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))

{konduktor_nightly-0.1.0.dev20250722105323 → konduktor_nightly-0.1.0.dev20250723105251}/konduktor/backends/jobset.py

@@ -16,7 +16,7 @@ if typing.TYPE_CHECKING:
 
 from konduktor import config, logging
 from konduktor.backends import backend, jobset_utils
-from konduktor.utils import kubernetes_utils, loki_utils, rich_utils, ux_utils
+from konduktor.utils import kubernetes_utils, log_utils, rich_utils, ux_utils
 
 Path = str
 logger = logging.get_logger(__file__)
@@ -193,7 +193,7 @@ class JobsetBackend(backend.Backend):
                     _wait_for_jobset_start(namespace, task.name)
                 try:
                     log_thread = threading.Thread(
-                        target=loki_utils.tail_loki_logs_ws,
+                        target=log_utils.tail_logs,
                         args=(task.name,),
                         daemon=True,
                     )

{konduktor_nightly-0.1.0.dev20250722105323 → konduktor_nightly-0.1.0.dev20250723105251}/konduktor/backends/jobset_utils.py

@@ -595,7 +595,7 @@ def show_status_table(namespace: str, all_users: bool):
             rows.append(
                 [
                     job['metadata']['name'],
-                    _get_status_string_colorized(job['status']),
+                    _get_status_string_colorized(job.get('status', {})),
                     _get_resources(job),
                     *_get_time_delta(job['metadata']['creationTimestamp']),
                 ]

{konduktor_nightly-0.1.0.dev20250722105323 → konduktor_nightly-0.1.0.dev20250723105251}/konduktor/cli.py

@@ -56,7 +56,6 @@ from konduktor.utils import (
     common_utils,
     kubernetes_utils,
     log_utils,
-    loki_utils,
     ux_utils,
 )
 
@@ -554,7 +553,8 @@ def status(all_users: bool):
     ),
 )
 @click.option(
-    '--tail',
+    '--num-lines',
+    '--num_lines' '-n',
     default=1000,
     type=int,
     help=(
@@ -568,10 +568,10 @@ def logs(
     status: bool,
     job_id: str,
     follow: bool,
-    tail: int,
+    num_lines: int,
 ):
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
-    """Tail the log of a job."""
+    """Retrieve/tail the log of a job."""
     if status:
         raise click.UsageError('`--status` is being deprecated)')
 
@@ -598,7 +598,7 @@ def logs(
         'Logs are tailed from 1 hour ago, ' 'to see more logs, check Grafana.',
         fg='yellow',
     )
-    loki_utils.tail_loki_logs_ws(job_id, follow=follow, num_logs=tail)
+    log_utils.tail_logs(job_id, follow=follow, num_logs=num_lines)
 
 
 @cli.command(cls=_DocumentedCodeCommand)

{konduktor_nightly-0.1.0.dev20250722105323 → konduktor_nightly-0.1.0.dev20250723105251}/konduktor/utils/log_utils.py

@@ -10,19 +10,41 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import copy
+import enum
 import io
+import json
 import multiprocessing
 import os
 import subprocess
 import sys
 import types
-from typing import List, Optional, Tuple, Type, Union
+import urllib.parse
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
+import colorama
+import kr8s
 import prettytable
+import requests
+import websockets
 
+from konduktor import config, logging
 from konduktor.utils import subprocess_utils
 
+logger = logging.get_logger(__name__)
+
+
+LOKI_REMOTE_PORT = 3100
+WEBSOCKET_TIMEOUT = 10
+INFINITY = 999999
+VICKY_REMOTE_PORT = 9428
+
+
+class LogBackend(enum.Enum):
+    VICTORIA = 'victoria'
+    LOKI = 'loki'
+
 
 class LineProcessor(object):
     """A processor for log lines."""
@@ -247,3 +269,147 @@ def run_with_log(
             # causing the stream handling stuck at `readline`.
             subprocess_utils.kill_children_processes()
             raise
+
+
+async def _read_logs(url: str, timeout: int, job_name: str, worker_id: int, port: int):
+    ws = await asyncio.wait_for(websockets.connect(url), timeout=WEBSOCKET_TIMEOUT)
+    logger.info(
+        f'{colorama.Fore.YELLOW}Tailing logs. '
+        f'Forwarding from remote port {port}. Press Ctrl+C to stop. '
+        f'{colorama.Style.RESET_ALL}'
+    )
+    try:
+        while True:
+            message = await asyncio.wait_for(ws.recv(), timeout=timeout)
+            try:
+                payload = json.loads(message)
+                for stream in payload['streams']:
+                    if stream['values'][0][1] is not None:
+                        print(
+                            f"{colorama.Fore.CYAN}{colorama.Style.BRIGHT} "
+                            f"(job_name={job_name} worker_id={worker_id})"
+                            f"{colorama.Style.RESET_ALL} {stream['values'][0][1]}",
+                            flush=True,
+                        )
+            except json.JSONDecodeError:
+                logger.warning(f'Failed to decode log skipping: {message}')
+                logger.debug(f'Dropped log: {message}')
+                continue
+    except asyncio.exceptions.TimeoutError:
+        logger.debug('Websocket timed-out, closing log stream!')
+    except KeyboardInterrupt:
+        logger.debug('Keyboard interrupt, closing log stream!')
+
+
+def tail_loki_logs_ws(
+    job_name: str, worker_id: int = 0, num_logs: int = 1000, follow: bool = True
+):
+    if num_logs > 5000:
+        # TODO(asaiacai): we should not have a limit on the number of logs, but rather
+        # let the user specify any number of lines, and we can print the last N lines.
+        # this can be done in chunks. Potentially, we can query range
+        # until we reach the end of the log and then invoke tail again.
+        # Also include checks that the job is running/ever ran.
+        raise ValueError('num_logs must be less than or equal to 5000')
+    loki_svc = kr8s.objects.Service.get('loki', namespace='loki')
+    with kr8s.portforward.PortForward(
+        loki_svc, LOKI_REMOTE_PORT, local_port='auto'
+    ) as port:
+        loki_url = f'ws://localhost:{port}/loki/api/v1/tail'
+        logger.debug(f'Loki URL: {loki_url}')
+        params = {
+            'query': urllib.parse.quote(
+                r'{' + f'k8s_job_name="{job_name}-workers-0",'
+                r' k8s_container_name="konduktor-container"} '
+                f' | batch_kubernetes_io_job_completion_index = `{worker_id}`'
+            ),
+            'limit': num_logs,
+            'delay': 5,
+            # TODO(asaiacai): need to auto-generate the start and end times.
+        }
+
+        query_string = '&'.join(f'{key}={value}' for key, value in params.items())
+        loki_url += f'?{query_string}'
+        timeout = INFINITY if follow else WEBSOCKET_TIMEOUT
+        asyncio.run(
+            _read_logs(loki_url, timeout, job_name, worker_id, LOKI_REMOTE_PORT)
+        )
+
+
+def tail_vicky_logs(
+    job_name: str, worker_id: int = 0, num_logs: int = 1000, follow: bool = True
+):
+    query: Dict[str, Any] = {}
+    if num_logs > 5000:
+        # TODO(asaiacai): we should not have a limit on the number of logs, but rather
+        # let the user specify any number of lines, and we can print the last N lines.
+        # this can be done in chunks. Potentially, we can query range
+        # until we reach the end of the log and then invoke tail again.
+        # Also include checks that the job is running/ever ran.
+        raise ValueError('num_logs must be less than or equal to 5000')
+    logger.info('ignoring num_logs argument for VictoriaLogs')
+    vicky_svc = kr8s.objects.Service.get(
+        'vls-victoria-logs-single-server', namespace='victoria-logs'
+    )
+    with kr8s.portforward.PortForward(
+        vicky_svc, VICKY_REMOTE_PORT, local_port='auto'
+    ) as port:
+        if follow:
+            timeout = INFINITY
+            vicky_url = f'http://localhost:{port}/select/logsql/tail'
+            query = {}
+        else:
+            vicky_url = f'http://localhost:{port}/select/logsql/query'
+            query = {'limit': num_logs}
+            timeout = 1
+        logger.debug(f'Vicky URL: {vicky_url}')
+
+        query['query'] = 'k8s.namespace.name: "default"'
+        query['start_offset'] = '1h'
+
+        try:
+            logger.debug(f'Making request to {vicky_url} with query: {query}')
+            with requests.post(
+                vicky_url, data=query, stream=True, timeout=timeout
+            ) as response:
+                logger.debug(f'Response status: {response.status_code}')
+                if response.status_code != 200:
+                    logger.error(
+                        f'VictoriaLogs API returned status {response.status_code}: '
+                        f'{response.text}'
+                    )
+                    return
+
+                for line in response.iter_lines(decode_unicode=True):
+                    if line:
+                        payload = json.loads(line)
+                        print(
+                            f"{colorama.Fore.CYAN}{colorama.Style.BRIGHT} "
+                            f"(job_name={job_name} worker_id={worker_id})"
+                            f"{colorama.Style.RESET_ALL} {payload['_msg']}",
+                            flush=True,
+                        )
+
+        except KeyboardInterrupt:
+            logger.info('\nStopping log stream...')
+        except requests.exceptions.Timeout:
+            logger.error(f'Request to VictoriaLogs timed out after {timeout} seconds')
+        except requests.exceptions.ConnectionError as e:
+            logger.error(f'Failed to connect to VictoriaLogs at {vicky_url}: {e}')
+        except requests.exceptions.RequestException as e:
+            logger.error(f'Request to VictoriaLogs failed: {e}')
+        except Exception as e:
+            logger.error(f'Unexpected error while tailing VictoriaLogs: {e}')
+
+
+def tail_logs(
+    job_name: str, worker_id: int = 0, num_logs: int = 1000, follow: bool = True
+):
+    logs_backend = config.get_nested(('logs', 'backend'), None)
+    if logs_backend == LogBackend.VICTORIA:
+        tail_vicky_logs(job_name, worker_id, num_logs, follow)
+    elif logs_backend == LogBackend.LOKI:
+        tail_loki_logs_ws(job_name, worker_id, num_logs, follow)
+    else:
+        logger.info('Defaulting to VictoriaLogs')
+        tail_vicky_logs(job_name, worker_id, num_logs, follow)

{konduktor_nightly-0.1.0.dev20250722105323 → konduktor_nightly-0.1.0.dev20250723105251}/konduktor/utils/schemas.py

@@ -490,6 +490,18 @@ def get_config_schema():
         },
     }
 
+    logs_configs = {
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'backend': {
+                'type': 'string',
+                'case_insensitive_enum': ['loki', 'victoria'],
+            },
+        },
+    }
+
     gpu_configs = {
         'type': 'object',
         'required': [],
@@ -537,6 +549,7 @@ def get_config_schema():
             'admin_policy': admin_policy_schema,
             'nvidia_gpus': gpu_configs,
             'allowed_clouds': allowed_clouds,
+            'logs': logs_configs,
             'tailscale': tailscale_configs,
             'ssh': ssh_configs,
             **cloud_configs,

{konduktor_nightly-0.1.0.dev20250722105323 → konduktor_nightly-0.1.0.dev20250723105251}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "konduktor-nightly"
-version = "0.1.0.dev20250722105323"
+version = "0.1.0.dev20250723105251"
 description = "GPU Cluster Health Management"
 packages = [
     {include = "konduktor"}