jerry-thomas 0.3.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

datapipeline/cli/commands/run.py

@@ -1,274 +1,260 @@
-import time
-from itertools import islice
-from pathlib import Path
-from typing import Iterator, List, Optional, Tuple, Union
-
+import json
 import logging
-
-from datapipeline.cli.visuals import visual_sources
-from datapipeline.config.dataset.dataset import FeatureDatasetConfig
-from datapipeline.config.dataset.feature import FeatureRecordConfig
+from pathlib import Path
+from typing import Optional
+
+from datapipeline.cli.commands.build import run_build_if_needed
+from datapipeline.cli.commands.run_config import (
+    RunEntry,
+    resolve_run_entries,
+)
+from datapipeline.cli.commands.serve_pipeline import serve_with_runtime
+from datapipeline.cli.visuals.runner import run_job
+from datapipeline.cli.visuals.sections import sections_from_path
+from datapipeline.config.context import resolve_run_profiles
 from datapipeline.config.dataset.loader import load_dataset
-from datapipeline.config.run import RunConfig, load_named_run_configs
-from datapipeline.domain.vector import Vector
-from datapipeline.pipeline.context import PipelineContext
-from datapipeline.pipeline.pipelines import build_vector_pipeline
-from datapipeline.pipeline.stages import post_process, split_stage
-from datapipeline.runtime import Runtime
-from datapipeline.services.bootstrap import bootstrap
-from datapipeline.cli.commands.writers import writer_factory, Writer
-from tqdm.contrib.logging import logging_redirect_tqdm
+from datapipeline.config.tasks import ServeOutputConfig
+from datapipeline.io.output import OutputResolutionError
+from datapipeline.pipeline.artifacts import StageDemand, required_artifacts_for
 
 logger = logging.getLogger(__name__)
 
 
-def _coerce_log_level(
-    value: Optional[Union[str, int]],
-    *,
-    default: int = logging.WARNING,
-) -> int:
-    if value is None:
-        return default
-    if isinstance(value, int):
-        return value
-    name = str(value).upper()
-    if name not in logging._nameToLevel:
-        raise ValueError(f"Unsupported log level: {value}")
-    return logging._nameToLevel[name]
-
-
-def _resolve_run_entries(project_path: Path, run_name: Optional[str]) -> List[Tuple[Optional[str], Optional[RunConfig]]]:
-    try:
-        entries = load_named_run_configs(project_path)
-    except FileNotFoundError:
-        entries = []
-    except Exception as exc:
-        logger.error("Failed to load run configs: %s", exc)
-        raise SystemExit(2) from exc
-
-    if entries:
-        if run_name:
-            entries = [entry for entry in entries if entry[0] == run_name]
-            if not entries:
-                logger.error("Unknown run config '%s'", run_name)
-                raise SystemExit(2)
-    else:
-        if run_name:
-            logger.error("Project does not define run configs.")
-            raise SystemExit(2)
-        entries = [(None, None)]
-    return entries
-
-
-def _iter_runtime_runs(
-    project_path: Path,
-    run_name: Optional[str],
-    keep_override: Optional[str],
-) -> Iterator[Tuple[int, int, Optional[str], Runtime]]:
-    run_entries = _resolve_run_entries(project_path, run_name)
-    total_runs = len(run_entries)
-    for idx, (entry_name, run_cfg) in enumerate(run_entries, start=1):
-        runtime = bootstrap(project_path)
-        if run_cfg is not None:
-            runtime.run = run_cfg
-            split_keep = getattr(runtime.split, "keep", None)
-            runtime.split_keep = run_cfg.keep or split_keep
-        if keep_override:
-            runtime.split_keep = keep_override
-        yield idx, total_runs, entry_name, runtime
-
-
-def _limit_items(items: Iterator[Tuple[object, object]], limit: Optional[int]) -> Iterator[Tuple[object, object]]:
-    if limit is None:
-        yield from items
-    else:
-        yield from islice(items, limit)
-
-
-def _throttle_vectors(vectors: Iterator[Tuple[object, Vector]], throttle_ms: Optional[float]) -> Iterator[Tuple[object, Vector]]:
-    if not throttle_ms or throttle_ms <= 0:
-        yield from vectors
-        return
-    delay = throttle_ms / 1000.0
-    for item in vectors:
-        yield item
-        time.sleep(delay)
-
-
-def _normalize(key: object, payload: object) -> dict:
-    return {
-        "key": list(key) if isinstance(key, tuple) else key,
-        "values": getattr(payload, "values", payload),
+def _profile_debug_payload(profile) -> dict[str, object]:
+    entry = profile.entry
+    payload: dict[str, object] = {
+        "label": profile.label,
+        "idx": profile.idx,
+        "total": profile.total,
+        "entry": {
+            "name": entry.name,
+            "path": str(entry.path) if entry.path else None,
+        },
+        "stage": profile.stage,
+        "limit": profile.limit,
+        "throttle_ms": profile.throttle_ms,
+        "log_level": {
+            "name": profile.log_decision.name,
+            "value": profile.log_decision.value,
+        },
+        "visuals": {
+            "provider": profile.visuals.visuals,
+            "progress": profile.visuals.progress,
+        },
+        "output": {
+            "transport": profile.output.transport,
+            "format": profile.output.format,
+            "payload": profile.output.payload,
+            "destination": str(profile.output.destination)
+            if profile.output.destination
+            else None,
+        },
     }
+    cfg = entry.config
+    if cfg is not None:
+        payload["run_config"] = cfg.model_dump(
+            exclude_unset=True, exclude_none=True)
+    return payload
 
 
-def _serve(
-    items: Iterator[Tuple[object, object]],
-    limit: Optional[int],
-    *,
-    writer: Writer,
-) -> int:
-    """Iterate, normalize, write, return count. Writers do only I/O."""
-    count = 0
-    try:
-        for key, payload in _limit_items(items, limit):
-            writer.write(_normalize(key, payload))
-            count += 1
-    except KeyboardInterrupt:
-        pass
-    finally:
-        writer.close()
-    return count
-
-
-def _report_end(output: Optional[str], count: int) -> None:
-    mode = (output or "print").lower()
-    if output and output.lower().endswith(".pt"):
-        logger.info("Saved %d items to %s", count, output)
-    elif output and output.lower().endswith(".csv"):
-        logger.info("Saved %d items to %s", count, output)
-    elif output and (output.lower().endswith(".jsonl.gz") or output.lower().endswith(".gz")):
-        logger.info("Saved %d items to %s", count, output)
-    elif mode == "stream":
-        logger.info("(streamed %d items)", count)
-    elif mode == "print":
-        logger.info("(printed %d items to stdout)", count)
-    else:
-        raise ValueError("unreachable: unknown output mode in _report_end")
-
+def _log_profile_start_debug(profile) -> None:
+    if not logger.isEnabledFor(logging.DEBUG):
+        return
+    payload = _profile_debug_payload(profile)
+    logger.debug(
+        "Run profile start (%s/%s):\n%s",
+        profile.idx,
+        profile.total,
+        json.dumps(payload, indent=2, default=str),
+    )
 
-def _serve_with_runtime(
-    runtime,
-    dataset: FeatureDatasetConfig,
-    limit: Optional[int],
-    output: Optional[str],
-    include_targets: bool,
-    throttle_ms: Optional[float],
-    stage: Optional[int] = None,
-) -> None:
-    context = PipelineContext(runtime)
 
-    features = list(dataset.features or [])
-    if include_targets:
-        features += list(dataset.targets or [])
+def _entry_sections(run_root: Optional[Path], entry: RunEntry) -> tuple[str, ...]:
+    # Prefix sections with a phase label for visuals; keep path-based detail.
+    path_sections = sections_from_path(run_root, entry.path)
+    return ("Run Tasks",) + tuple(path_sections[1:])
 
-    if not features:
-        logger.warning("(no features configured; nothing to serve)")
-        return
 
-    if stage is not None and stage <= 5:
-        for cfg in features:
-            stream = build_vector_pipeline(
-                context,
-                [cfg],
-                dataset.group_by,
-                stage=stage,
-            )
-            items = ((cfg.id, item) for item in stream)
-            writer = writer_factory(output)
-            count = _serve(items, limit, writer=writer)
-            _report_end(output, count)
-        return
+def _build_cli_output_config(
+    transport: Optional[str],
+    fmt: Optional[str],
+    path: Optional[str],
+    payload: Optional[str],
+) -> tuple[ServeOutputConfig | None, Optional[str]]:
+    payload_style = None
+    if payload is not None:
+        payload_style = payload.lower()
+        if payload_style not in {"sample", "vector"}:
+            logger.error("--out-payload must be 'sample' or 'vector'")
+            raise SystemExit(2)
 
-    vector_stage = 6 if stage in (6, 7) else None
-    vectors = build_vector_pipeline(
-        context,
-        features,
-        dataset.group_by,
-        stage=vector_stage,
+    if transport is None and fmt is None and path is None:
+        return None, payload_style
+
+    if not transport or not fmt:
+        logger.error(
+            "--out-transport and --out-format must be provided together")
+        raise SystemExit(2)
+    transport = transport.lower()
+    fmt = fmt.lower()
+    if transport == "fs":
+        if not path:
+            logger.error(
+                "--out-path is required when --out-transport=fs (directory)")
+            raise SystemExit(2)
+        return (
+            ServeOutputConfig(
+                transport="fs",
+                format=fmt,
+                directory=Path(path),
+                payload=payload_style or "sample",
+            ),
+            None,
+        )
+    if path:
+        logger.error("--out-path is only valid when --out-transport=fs")
+        raise SystemExit(2)
+    return (
+        ServeOutputConfig(
+            transport="stdout",
+            format=fmt,
+            payload=payload_style or "sample",
+        ),
+        None,
     )
 
-    if stage in (None, 7):
-        vectors = post_process(context, vectors)
-    if stage is None:
-        vectors = split_stage(runtime, vectors)
-        vectors = _throttle_vectors(vectors, throttle_ms)
-
-    writer = writer_factory(output)
-    result_count = _serve(vectors, limit, writer=writer)
-    _report_end(output, result_count)
 
-
-def _execute_runs(
+def ensure_stage_artifacts(
     project_path: Path,
-    stage: Optional[int],
-    limit: Optional[int],
-    output: Optional[str],
-    include_targets: Optional[bool],
-    keep: Optional[str],
-    run_name: Optional[str],
+    dataset,
+    profiles,
     *,
-    cli_log_level: Optional[str],
-    base_log_level: str,
+    cli_visuals: Optional[str],
+    cli_progress: Optional[str],
+    workspace,
 ) -> None:
-    # Helper for precedence: CLI > config > default
-    def pick(cli_val, cfg_val, default=None):
-        return cli_val if cli_val is not None else (cfg_val if cfg_val is not None else default)
-
-    dataset_name = "vectors" if stage is None else "features"
-    dataset = load_dataset(project_path, dataset_name)
-
-    base_level_name = str(base_log_level).upper()
-    base_level_value = _coerce_log_level(base_level_name)
-
-    for idx, total_runs, entry_name, runtime in _iter_runtime_runs(project_path, run_name, keep):
-        run = getattr(runtime, "run", None)
-
-        # resolving argument hierarchy CLI args > run config > defaults
-        resolved_limit = pick(limit, getattr(run, "limit", None), None)
-        resolved_output = pick(output, getattr(run, "output", None), "print")
-        resolved_include_targets = pick(
-            include_targets, getattr(run, "include_targets", None), False)
-        throttle_ms = getattr(run, "throttle_ms", None)
-        resolved_level_name = pick(
-            cli_log_level.upper() if cli_log_level else None,
-            getattr(run, "log_level", None),
-            base_level_name,
-        )
-        resolved_level_value = _coerce_log_level(
-            resolved_level_name, default=base_level_value)
-
-        root_logger = logging.getLogger()
-        if root_logger.level != resolved_level_value:
-            root_logger.setLevel(resolved_level_value)
-
-        label = entry_name or f"run{idx}"
-        logger.info("Run '%s' (%d/%d)", label, idx, total_runs)
-
-        with visual_sources(runtime, resolved_level_value):
-            with logging_redirect_tqdm():
-                _serve_with_runtime(
-                    runtime,
-                    dataset,
-                    limit=resolved_limit,
-                    output=resolved_output,
-                    include_targets=resolved_include_targets,
-                    throttle_ms=throttle_ms,
-                    stage=stage,
-                )
+    demands = [StageDemand(profile.stage) for profile in profiles]
+    required = required_artifacts_for(dataset, demands)
+    if not required:
+        return
+    run_build_if_needed(
+        project_path,
+        cli_visuals=cli_visuals,
+        cli_progress=cli_progress,
+        workspace=workspace,
+        required_artifacts=required,
+    )
 
 
 def handle_serve(
     project: str,
     limit: Optional[int],
-    output: Optional[str],
-    include_targets: Optional[bool] = None,
     keep: Optional[str] = None,
     run_name: Optional[str] = None,
     stage: Optional[int] = None,
+    out_transport: Optional[str] = None,
+    out_format: Optional[str] = None,
+    out_payload: Optional[str] = None,
+    out_path: Optional[str] = None,
+    skip_build: bool = False,
     *,
     cli_log_level: Optional[str],
     base_log_level: str,
+    cli_visuals: Optional[str] = None,
+    cli_progress: Optional[str] = None,
+    workspace=None,
 ) -> None:
     project_path = Path(project)
-    _execute_runs(
-        project_path=project_path,
-        stage=stage,
-        limit=limit,
-        output=output,
-        include_targets=include_targets,
-        keep=keep,
-        run_name=run_name,
-        cli_log_level=cli_log_level,
-        base_log_level=base_log_level,
-    )
+    run_entries, run_root = resolve_run_entries(project_path, run_name)
+
+    cli_output_cfg, payload_override = _build_cli_output_config(
+        out_transport, out_format, out_path, out_payload)
+    try:
+        profiles = resolve_run_profiles(
+            project_path=project_path,
+            run_entries=run_entries,
+            keep=keep,
+            stage=stage,
+            limit=limit,
+            cli_output=cli_output_cfg,
+            cli_payload=payload_override or (
+                out_payload.lower() if out_payload else None),
+            workspace=workspace,
+            cli_log_level=cli_log_level,
+            base_log_level=base_log_level,
+            cli_visuals=cli_visuals,
+            cli_progress=cli_progress,
+            create_run=False,
+        )
+    except OutputResolutionError as exc:
+        logger.error("Invalid output configuration: %s", exc)
+        raise SystemExit(2) from exc
+
+    vector_dataset = load_dataset(project_path, "vectors")
+    skip_reason = None
+    if skip_build:
+        skip_reason = "--skip-build flag provided"
+
+    if not skip_reason:
+        ensure_stage_artifacts(
+            project_path,
+            vector_dataset,
+            profiles,
+            cli_visuals=cli_visuals,
+            cli_progress=cli_progress,
+            workspace=workspace,
+        )
+        profiles = resolve_run_profiles(
+            project_path=project_path,
+            run_entries=run_entries,
+            keep=keep,
+            stage=stage,
+            limit=limit,
+            cli_output=cli_output_cfg,
+            cli_payload=payload_override or (
+                out_payload.lower() if out_payload else None),
+            workspace=workspace,
+            cli_log_level=cli_log_level,
+            base_log_level=base_log_level,
+            cli_visuals=cli_visuals,
+            cli_progress=cli_progress,
+            create_run=True,
+        )
+
+    datasets: dict[str, object] = {}
+    datasets["vectors"] = vector_dataset
+    for profile in profiles:
+        dataset_name = "vectors" if profile.stage is None else "features"
+        dataset = datasets.get(dataset_name)
+        if dataset is None:
+            dataset = load_dataset(project_path, dataset_name)
+            datasets[dataset_name] = dataset
+
+        root_logger = logging.getLogger()
+        if root_logger.level != profile.log_decision.value:
+            root_logger.setLevel(profile.log_decision.value)
+
+        def _work(profile=profile):
+            _log_profile_start_debug(profile)
+            serve_with_runtime(
+                profile.runtime,
+                dataset,
+                limit=profile.limit,
+                target=profile.output,
+                throttle_ms=profile.throttle_ms,
+                stage=profile.stage,
+                visuals=profile.visuals.visuals,
+            )
+
+        sections = _entry_sections(run_root, profile.entry)
+        run_job(
+            sections=sections,
+            label=profile.label,
+            visuals=profile.visuals.visuals or "auto",
+            progress_style=profile.visuals.progress or "auto",
+            level=profile.log_decision.value,
+            runtime=profile.runtime,
+            work=_work,
+            idx=profile.idx,
+            total=profile.total,
+        )

datapipeline/cli/commands/run_config.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Iterator, List, NamedTuple, Optional, Sequence
+
+from datapipeline.config.tasks import ServeTask, serve_tasks
+from datapipeline.runtime import Runtime
+from datapipeline.services.bootstrap import bootstrap
+
+logger = logging.getLogger(__name__)
+
+
+class RunEntry(NamedTuple):
+    name: Optional[str]
+    config: Optional[ServeTask]
+    path: Optional[Path]
+
+
+def resolve_run_entries(project_path: Path, run_name: Optional[str]) -> tuple[List[RunEntry], Optional[Path]]:
+    try:
+        raw_entries = serve_tasks(project_path)
+    except FileNotFoundError:
+        raw_entries = []
+    except Exception as exc:
+        logger.error("Failed to load serve tasks: %s", exc)
+        raise SystemExit(2) from exc
+
+    entries: List[RunEntry] = []
+    root_path: Optional[Path] = None
+
+    if raw_entries:
+        if not run_name:
+            raw_entries = [task for task in raw_entries if task.enabled]
+        if run_name:
+            raw_entries = [
+                task
+                for task in raw_entries
+                if task.effective_name() == run_name
+            ]
+            if not raw_entries:
+                logger.error("Unknown run task '%s'", run_name)
+                raise SystemExit(2)
+        for task in raw_entries:
+            path = getattr(task, "source_path", None)
+            if root_path is None and path is not None:
+                root_path = path.parent
+            entries.append(
+                RunEntry(
+                    name=task.effective_name(),
+                    config=task,
+                    path=path,
+                )
+            )
+    else:
+        if run_name:
+            logger.error("Project does not define serve tasks.")
+            raise SystemExit(2)
+        entries = [RunEntry(name=None, config=None, path=None)]
+    return entries, root_path
+
+
+def iter_runtime_runs(
+    project_path: Path,
+    run_entries: Sequence[RunEntry],
+    keep_override: Optional[str],
+) -> Iterator[tuple[int, int, RunEntry, Runtime]]:
+    total_runs = len(run_entries)
+    for idx, entry in enumerate(run_entries, start=1):
+        run_cfg = entry.config
+        runtime = bootstrap(project_path)
+        if run_cfg is not None:
+            runtime.run = run_cfg
+            split_keep = getattr(runtime.split, "keep", None)
+            runtime.split_keep = run_cfg.keep or split_keep
+        if keep_override:
+            runtime.split_keep = keep_override
+        yield idx, total_runs, entry, runtime
+
+
+def determine_preview_stage(
+    cli_stage: Optional[int],
+    run_entries: Sequence[RunEntry],
+) -> tuple[Optional[int], Optional[str]]:
+    if cli_stage is not None:
+        return cli_stage, "CLI flag"
+
+    stages: List[int] = []
+    for entry in run_entries:
+        run_cfg = entry.config
+        cfg_stage = getattr(run_cfg, "stage", None) if run_cfg else None
+        if cfg_stage is None:
+            return None, None
+        stages.append(cfg_stage)
+
+    if not stages or any(stage > 5 for stage in stages):
+        return None, None
+
+    if len(set(stages)) == 1:
+        return stages[0], "run config"
+    return min(stages), "run configs"