daggerml.api

   1from __future__ import annotations
   2
   3import logging
   4import time
   5import traceback
   6from collections.abc import Iterator, Mapping, Sequence
   7from contextlib import contextmanager
   8from contextvars import ContextVar
   9from dataclasses import dataclass, field
  10from importlib import metadata
  11from tempfile import TemporaryDirectory
  12from threading import RLock
  13from typing import Any, Optional, Protocol, Union, cast, get_args, overload
  14
  15from daggerml._core import CancellationError, Dml, DmlRepoError, Error, Ref, Runnable, Uri
  16from daggerml.util import BackoffWithJitter, current_time_millis
  17
  18logger = logging.getLogger(__name__)
  19
  20Scalar = Union[str, int, float, bool, type(None), Uri, Runnable]
  21Collection = Union[list, dict]
  22ProjectionStep = Union[str, int, list[int]]
  23
  24_NO_DEFAULT_DML = object()
  25_SCOPED_DEFAULT_DML: ContextVar[object] = ContextVar("daggerml_scoped_default_dml", default=_NO_DEFAULT_DML)
  26_PROCESS_DEFAULT_DML: Optional["Dml"] = None
  27
  28
  29def _resolve_default_dml(*, create: bool = True) -> tuple["Dml", str]:
  30    scoped = _SCOPED_DEFAULT_DML.get()
  31    if scoped is not _NO_DEFAULT_DML:
  32        return cast("Dml", scoped), "scoped"
  33
  34    global _PROCESS_DEFAULT_DML
  35    if _PROCESS_DEFAULT_DML is not None:
  36        return _PROCESS_DEFAULT_DML, "process"
  37
  38    if not create:
  39        raise DmlRepoError("No default Dml is configured")
  40
  41    _PROCESS_DEFAULT_DML = Dml()
  42    return _PROCESS_DEFAULT_DML, "implicit"
  43
  44
  45def get_default_dml() -> "Dml":
  46    """Return the active default Dml runtime."""
  47    dml, _source = _resolve_default_dml(create=True)
  48    return dml
  49
  50
  51def set_default_dml(dml: "Dml") -> None:
  52    """Set the process-default Dml runtime."""
  53    global _PROCESS_DEFAULT_DML
  54    _PROCESS_DEFAULT_DML = dml
  55
  56
  57def clear_default_dml() -> None:
  58    """Clear the process-default Dml runtime."""
  59    global _PROCESS_DEFAULT_DML
  60    _PROCESS_DEFAULT_DML = None
  61
  62
  63@contextmanager
  64def use_default_dml(dml: "Dml"):
  65    """Temporarily override the default Dml runtime for the active context."""
  66    token = _SCOPED_DEFAULT_DML.set(dml)
  67    try:
  68        yield dml
  69    finally:
  70        _SCOPED_DEFAULT_DML.reset(token)
  71
  72
  73def new(
  74    name="",
  75    message="",
  76    cache_key: str | None = None,
  77    execution_id: str | None = None,
  78    tags: list[str] | None = None,
  79    dml: Dml | None = None,
  80) -> "Dag":
  81    """Create a new DAG using the active or provided Dml runtime."""
  82    runtime = dml or get_default_dml()
  83    execution = Ref(f"index:{execution_id}") if execution_id is not None else None
  84    create_kwargs = {"cache_key": cache_key, "execution": execution}
  85    if tags is not None:
  86        create_kwargs["tags"] = tags
  87    index_id = runtime.runtime.create(**create_kwargs)
  88    return Dag(dml=runtime, token=index_id, name=name, message=message)
  89
  90
  91def load(
  92    name: str,
  93    dml: Dml | None = None,
  94    *,
  95    revision: Ref | str = "HEAD",
  96    remote: bool = False,
  97    dep: str | None = None,
  98) -> "Dag":
  99    """Load a DAG using the active default Dml runtime."""
 100    dml = dml or get_default_dml()
 101    dag_ref = dml.show(revision, remote=remote, dep=dep)["dags"].get(name)
 102    if dag_ref is None:
 103        raise DmlRepoError(f"DAG not found: {name}")
 104    return Dag(dml=dml, ref=dag_ref, name=name)
 105
 106
 107def resume(
 108    frozen: Ref,
 109    *,
 110    name: str,
 111    message: str,
 112    dml: Dml | None = None,
 113) -> "Dag":
 114    """Resume a frozen DAG runtime with explicitly supplied commit metadata."""
 115    runtime = dml or get_default_dml()
 116    index_id = runtime.runtime.unfreeze(frozen)
 117    return Dag(dml=runtime, token=index_id, name=name, message=message)
 118
 119
 120@contextmanager
 121def temporary(prefix="dml-tmp-", **kw):
 122    """Create a temporary Dml runtime with an unborn attached HEAD."""
 123    with TemporaryDirectory(prefix=prefix) as tmpdir:
 124        yield Dml.init(project_home=tmpdir, **kw)
 125
 126
 127def status() -> dict[str, object]:
 128    """Return status for the active default Dml runtime."""
 129    dml, source = _resolve_default_dml(create=True)
 130    return {
 131        "default": {
 132            "source": source,
 133            "has_scoped_override": _SCOPED_DEFAULT_DML.get() is not _NO_DEFAULT_DML,
 134            "has_process_default": _PROCESS_DEFAULT_DML is not None,
 135        },
 136        "status": dml.status(),
 137    }
 138
 139
 140def _make_node(dag: "Dag", ref: Ref) -> "Node":
 141    """
 142    Create a Node from a Dag and Ref.
 143
 144    Parameters
 145    ----------
 146    dag : Dag
 147        The parent DAG.
 148    ref : Ref
 149        The reference to the node.
 150    Returns
 151    -------
 152    Node
 153        A Node instance representing the reference in the DAG.
 154    """
 155    node_value = dag.dml.dag.get_node(ref)
 156    if isinstance(node_value, Error):
 157        raise NodeError(node_value, node_ref=ref, dag=dag)
 158    info: dict[str, Any] = {"data_type": type(node_value).__name__.lower()}
 159    # Determine node type based on value and populate info
 160    if isinstance(node_value, list):
 161        info["length"] = len(node_value)
 162        node = ListNode(dag, ref, _info=info)
 163    elif isinstance(node_value, dict):
 164        info["length"] = len(node_value)
 165        info["keys"] = sorted(node_value.keys())
 166        node = DictNode(dag, ref, _info=info)
 167    elif isinstance(node_value, Runnable):
 168        node = RunnableNode(dag, ref, _info=info)
 169    else:
 170        node = ScalarNode(dag, ref, _info=info)
 171    return node
 172
 173
 174def _info_for_value(value: Any) -> dict[str, Any]:
 175    info: dict[str, Any] = {"data_type": type(value).__name__.lower()}
 176    if isinstance(value, list):
 177        info["length"] = len(value)
 178    elif isinstance(value, dict):
 179        info["length"] = len(value)
 180        info["keys"] = sorted(value.keys())
 181    return info
 182
 183
 184def _normalize_projection_step(key: ProjectionStep | slice, *, length: int | None = None) -> ProjectionStep:
 185    if isinstance(key, slice):
 186        if key.step is not None:
 187            raise ValueError("Slice step is not supported")
 188        start = key.start if key.start is not None else 0
 189        if key.stop is None:
 190            if length is None:
 191                raise DmlRepoError("Slice stop requires known collection length")
 192            stop = length
 193        else:
 194            stop = key.stop
 195        return [start, stop]
 196    return key
 197
 198
 199def _apply_projection_step(value: Any, step: ProjectionStep) -> Any:
 200    if isinstance(step, list):
 201        if len(step) != 2:
 202            raise DmlRepoError("Slice projection requires exactly [start, stop]")
 203        return value[slice(*step)]
 204    return value[step]
 205
 206
 207def _apply_projection_path(value: Any, path: tuple[ProjectionStep, ...]) -> Any:
 208    for step in path:
 209        value = _apply_projection_step(value, step)
 210    return value
 211
 212
 213def _describe_node(node: "Node") -> Mapping[str, Any]:
 214    return node.dag.dml.dag.describe_node(node.ref)
 215
 216
 217def _builtin_name_for_argv(dag: "Dag", argv_refs: list[Ref]) -> str | None:
 218    if not argv_refs:
 219        return None
 220    runnable = dag.dml.dag.get_node(argv_refs[0], recursive=True)
 221    if not isinstance(runnable, Runnable):
 222        return None
 223    if runnable.adapter != "":
 224        return None
 225    if not runnable.target.uri.startswith("daggerml:"):
 226        return None
 227    return runnable.target.uri.split(":", 1)[1]
 228
 229
 230def _prepend_get_path_step(path: tuple[ProjectionStep, ...], key: ProjectionStep) -> tuple[ProjectionStep, ...] | None:
 231    if isinstance(key, list):
 232        if len(key) != 2:
 233            return None
 234        if not path:
 235            return None
 236        first, *rest = path
 237        if not isinstance(first, int):
 238            return None
 239        return (key[0] + first, *rest)
 240    return (key, *path)
 241
 242
 243def _backtrack_builtin(
 244    node: "Node", argv_refs: list[Ref], path: tuple[ProjectionStep, ...]
 245) -> tuple["Node", tuple[ProjectionStep, ...]] | None:
 246    builtin = _builtin_name_for_argv(node.dag, argv_refs)
 247    if builtin is None:
 248        return None
 249    arg_nodes = [_make_node(node.dag, ref) for ref in argv_refs[1:]]
 250    if builtin == "get":
 251        if len(arg_nodes) < 2:
 252            return None
 253        key = cast(ProjectionStep, arg_nodes[1].value())
 254        next_path = _prepend_get_path_step(path, key)
 255        if next_path is None:
 256            return None
 257        return arg_nodes[0], next_path
 258    if builtin == "list":
 259        if not path:
 260            return None
 261        index, *rest = path
 262        if not isinstance(index, int):
 263            return None
 264        if index < 0:
 265            index += len(arg_nodes)
 266        if index < 0 or index >= len(arg_nodes):
 267            return None
 268        return arg_nodes[index], tuple(rest)
 269    if builtin == "dict":
 270        if not path:
 271            return None
 272        key, *rest = path
 273        if not isinstance(key, str):
 274            return None
 275        for idx in range(0, len(arg_nodes), 2):
 276            if idx + 1 >= len(arg_nodes):
 277                break
 278            if arg_nodes[idx].value() == key:
 279                return arg_nodes[idx + 1], tuple(rest)
 280        return None
 281    if builtin == "assoc":
 282        if len(arg_nodes) < 3 or not path:
 283            return None
 284        selected_key, *rest = path
 285        assoc_key = arg_nodes[1].value()
 286        if selected_key == assoc_key:
 287            return arg_nodes[2], tuple(rest)
 288        return arg_nodes[0], path
 289    if builtin == "conj":
 290        if len(arg_nodes) < 2 or not path:
 291            return None
 292        index, *rest = path
 293        if not isinstance(index, int):
 294            return None
 295        base_len = len(arg_nodes[0].value())
 296        if index < 0:
 297            index += base_len + 1
 298        if index == base_len:
 299            return arg_nodes[1], tuple(rest)
 300        return arg_nodes[0], path
 301    return None
 302
 303
 304def _nearest_context_state(
 305    node: "Node", path: tuple[ProjectionStep, ...]
 306) -> tuple[Dag, Node, tuple[ProjectionStep, ...], bool]:
 307    current = node
 308    current_path = path
 309    while True:
 310        node_info = _describe_node(current)
 311        node_type = node_info["type"]
 312        if node_type == "ImportNode":
 313            source_dag = Dag(dml=current.dag.dml, ref=cast(Ref, node_info["dag"]))
 314            source_node = _make_node(source_dag, cast(Ref, node_info["node"]))
 315            if current_path:
 316                current = source_node
 317                continue
 318            return source_dag, source_node, current_path, True
 319        if node_type == "FnNode":
 320            argv_refs = cast(list[Ref], node_info["argv"])
 321            source = _backtrack_builtin(current, argv_refs, current_path)
 322            if source is not None:
 323                current, current_path = source
 324                continue
 325            fn_dag = Dag(dml=current.dag.dml, ref=cast(Ref, node_info["dag"]))
 326            return fn_dag, fn_dag.result, current_path, True
 327        return current.dag, current, current_path, False
 328
 329
 330def _resolve_context(node: "Node", path: tuple[ProjectionStep, ...], *, root: bool) -> Dag:
 331    context_dag, next_node, next_path, can_recurse = _nearest_context_state(node, path)
 332    if not root:
 333        return context_dag
 334    if not can_recurse:
 335        return context_dag
 336    return _resolve_context(next_node, next_path, root=True)
 337
 338
 339@dataclass
 340class Dag:
 341    dml: Dml
 342    token: Optional[Ref] = None  # Working index id
 343    ref: Optional[Ref] = None
 344    name: str = ""  # DAG name for commit
 345    message: str = ""  # Commit message
 346
 347    def __repr__(self):
 348        to = self.ref.to if self.ref else (self.token.to if self.token is not None else "NA")
 349        return f"Dag({to})"
 350
 351    def __hash__(self):
 352        "Useful only for tests."
 353        return 42
 354
 355    def __eq__(self, other):
 356        "DAG equality is based on identity, not content."
 357        if not isinstance(other, Dag):
 358            return False
 359        return self.ref == other.ref and self.token == other.token and self.dml == other.dml
 360
 361    def __enter__(self):
 362        "Catch exceptions and commit an Error"
 363        assert not self.ref
 364        return self
 365
 366    def __exit__(self, exc_type, exc_value, tb):
 367        if exc_value is not None:
 368            # Convert exception to Error and commit it
 369            traceback.print_exception(exc_type, exc_value, tb)
 370            err = Error.from_ex(exc_value) if not isinstance(exc_value, Error) else exc_value
 371            self.commit(err)
 372
 373    def _require_index_ref(self) -> Ref:
 374        if self.token is None:
 375            raise DmlRepoError("No active index")
 376        return self.token
 377
 378    def _read_dag_ref(self) -> Ref:
 379        """Return the completed or partial DAG backing this wrapper's reads."""
 380        if self.ref is not None:
 381            return self.ref
 382        return cast(Ref, self.dml.runtime.describe(self._require_index_ref())["dag"])
 383
 384    @property
 385    def tags(self) -> list[str]:
 386        """Return the normalized tags stored on this DAG."""
 387        return list(self.dml.dag.describe(self._read_dag_ref())["tags"])
 388
 389    def _put_literal(self, value: Any, *, name: Optional[str] = None) -> Ref:
 390        index_id = self._require_index_ref()
 391        value = apply_codecs(value, dag=self)
 392        return self.dml.runtime.put_literal(index_id, value, name=name)
 393
 394    def _start_fn(self, argv: list[Ref], *, name: Optional[str] = None) -> Optional[Ref]:
 395        return self.dml.runtime.start_fn(self._require_index_ref(), argv, name=name)
 396
 397    def _call_builtin(self, uri: str, *args: Any, name: Optional[str] = None) -> Ref:
 398        fn_ref = self._put_literal(Runnable(target=Uri(uri), kwargs={}, adapter=""))
 399        argv: list[Ref] = [fn_ref]
 400        for arg in args:
 401            argv.append(arg if isinstance(arg, Ref) else self._put_literal(arg))
 402        result = self._start_fn(argv, name=name)
 403        if result is None:
 404            raise DmlRepoError("Function execution failed")
 405        return result
 406
 407    def __len__(self) -> int:
 408        return len(self.keys())
 409
 410    def __iter__(self):
 411        yield from self.keys()
 412
 413    def _get_named_node(self, name: str) -> "Node":
 414        node_ref = self.dml.dag.describe(self._read_dag_ref())["names"].get(name)
 415        if node_ref is None:
 416            raise DmlRepoError(f"Node '{name}' not found in DAG")
 417        return _make_node(self, node_ref)
 418
 419    def _set_named_node(self, name: str, value: Any) -> None:
 420        if self.ref is not None:
 421            raise DmlRepoError("Cannot set node names on a committed DAG.")
 422        if isinstance(value, Node):
 423            value = value.ref
 424        if isinstance(value, Ref):
 425            self.dml.runtime.set_node_name(self._require_index_ref(), name, value)
 426            return
 427        self.put(value, name=name)
 428
 429    def __getitem__(self, name: str) -> "Node":
 430        return self._get_named_node(name)
 431
 432    def __setitem__(self, name: str, value: Any) -> None:
 433        self.put(value, name=name)
 434
 435    def __getattr__(self, name: str) -> "Node":
 436        if name.startswith("_"):
 437            raise AttributeError(name)
 438        return self[name]
 439
 440    def __setattr__(self, name: str, value: Any) -> None:
 441        dataclass_fields = getattr(type(self), "__dataclass_fields__", {})
 442        if name in dataclass_fields or hasattr(type(self), name):
 443            object.__setattr__(self, name, value)
 444            return
 445        self[name] = value
 446
 447    def keys(self) -> list[str]:
 448        """Get the list of all node names in the dag"""
 449        names_dict = self.dml.dag.describe(self._read_dag_ref())["names"]
 450        return sorted(names_dict.keys())
 451
 452    def values(self) -> list["Node"]:
 453        """Get the list of all nodes in the dag"""
 454        names_dict = self.dml.dag.describe(self._read_dag_ref())["names"]
 455        return [_make_node(self, ref) for ref in names_dict.values()]
 456
 457    @property
 458    def argv(self) -> "ListNode":
 459        "Access the dag's argv node"
 460        argv_ref = self.dml.dag.describe(self._read_dag_ref())["argv"]
 461        if not isinstance(argv_ref, Ref):
 462            raise DmlRepoError(f"'{self.__class__.__name__}' dag has no argv")
 463        return cast(ListNode, _make_node(self, argv_ref))
 464
 465    @property
 466    def result(self) -> "Node":
 467        """Get the result node of the dag"""
 468        if self.ref is None:
 469            raise DmlRepoError("Cannot access result of an uncommitted DAG")
 470        description = self.dml.dag.describe(self.ref)
 471        error_ref = description.get("error")
 472        if isinstance(error_ref, Ref):
 473            raise self.dml.dag.get_error(error_ref)
 474        ref = description.get("result")
 475        if not isinstance(ref, Ref):
 476            raise DmlRepoError(f"'{self.__class__.__name__}' dag has not been committed yet")
 477        return _make_node(self, ref)
 478
 479    @overload
 480    def put(self, value: Union[list, "ListNode"], *, name=None) -> "ListNode": ...
 481    @overload
 482    def put(self, value: Union[dict, "DictNode"], *, name=None) -> "DictNode": ...
 483    @overload
 484    def put(self, value: Union[Runnable, "RunnableNode"], *, name=None) -> "RunnableNode": ...
 485    @overload
 486    def put(self, value: Union[Scalar, "ScalarNode"], *, name=None) -> "ScalarNode": ...
 487    @overload
 488    def put(self, value: Any, *, name=None) -> "Node": ...
 489    def put(self, value: Any, *, name=None) -> "Node":
 490        """
 491        Add a value to the DAG.
 492
 493        Parameters
 494        ----------
 495        value : Union[Scalar, Collection]
 496            Value to add
 497        name : str, optional
 498            Name for the node
 499        Returns
 500        -------
 501        Node
 502            Node representing the value
 503
 504        Examples
 505        --------
 506        >>> n1 = dag.put(42, name="answer")
 507        >>> n1.value()
 508        42
 509        >>> n2 = dag.put({"a": 1, "b": [n1, "23"]})
 510        >>> n2.value()
 511        {'a': 1, 'b': [42, '23']}
 512        >>> n3 = dag.put({"a": 1, "b": [n1, "23"]})
 513        >>> n3.value()
 514        {'a': 1, 'b': [42, '23']}
 515        """
 516        return _make_node(self, self._put_literal(value, name=name))
 517
 518    def require(self, dag_name: str | "Dag", node_name: str | None = None, *, name: str | None = None) -> "Node":
 519        """
 520        Import a node from a different (committed) DAG into the current DAG.
 521
 522        Parameters
 523        ----------
 524        dag_name : str
 525            Name of the DAG to import from
 526        node_name : str, optional
 527            Name of the node to import. If None, imports the result node of the DAG.
 528
 529        Returns
 530        -------
 531        Node
 532            The loaded node or DAG
 533
 534        Examples
 535        --------
 536        >>> dag = new(dml=dml, name="test", message="test")
 537        >>> n1 = dag.put(42, name="answer")
 538        >>> n2 = dag.put({"a": 1, "b": [n1, "23"]}, name="data")
 539        >>> dag.commit(n2)
 540        >>> dag2 = new(dml=dml, name="test2", message="test2")
 541        >>> imported_n2 = dag2.require("test", "data", name="imported_data")
 542        >>> imported_n2.value()
 543        {'a': 1, 'b': [42, '23']}
 544        """
 545        index = self._require_index_ref()
 546        if isinstance(dag_name, Dag):
 547            if dag_name.ref is None:
 548                raise DmlRepoError("Cannot import an uncommitted DAG")
 549            dag_ref = dag_name.ref
 550        else:
 551            commit = self.dml.runtime.describe(index)["parents"][0]
 552            dag_ref = self.dml.show(revision=commit)["dags"].get(dag_name)
 553            if dag_ref is None:
 554                raise DmlRepoError(f"DAG not found: {dag_name}")
 555        dag_info = self.dml.dag.describe(dag_ref)
 556        node_ref = dag_info["names"].get(node_name) if node_name else dag_info.get("result")
 557        if node_ref is None:
 558            raise DmlRepoError(f"Node '{node_name}' not found in DAG '{dag_name}'")
 559        node_ref = self.dml.runtime.put_import(index, dag_ref, node_ref, name=name)
 560        return _make_node(self, node_ref)
 561
 562    def call(
 563        self,
 564        fn: Any,
 565        *args: Any,
 566        name: Optional[str] = None,
 567        sleep: Optional[callable] = None,
 568        timeout: int = -1,
 569    ) -> "Node":
 570        """
 571        Call a function node with arguments.
 572
 573        Parameters
 574        ----------
 575        fn : Union[Runnable, RunnableNode]
 576            Function to call
 577        *args : Union[Node, Scalar, Collection]
 578            Arguments to pass to the function
 579        name : str, optional
 580            Name for the result node
 581        sleep : callable, optional
 582            A nullary function that returns sleep time in milliseconds
 583        timeout : int, default=-1
 584            Maximum time to wait in milliseconds. If <= 0, wait indefinitely.
 585
 586        Returns
 587        -------
 588        Node
 589            Result node
 590
 591        Raises
 592        ------
 593        TimeoutError
 594            If the function call exceeds the timeout
 595        Error
 596            If the function returns an error
 597        """
 598        sleep = sleep or BackoffWithJitter()
 599        argv_seed = [fn, *args]
 600        end = current_time_millis() + timeout
 601        while timeout <= 0 or current_time_millis() < end:
 602            argv_refs = [self._put_literal(value) for value in argv_seed]
 603            resp = self._start_fn(argv_refs, name=name)
 604            if resp:
 605                return _make_node(self, resp)
 606            time.sleep(sleep() / 1000)
 607        raise TimeoutError(f"invoking function: {fn}")
 608
 609    def commit(self, value) -> None:
 610        """
 611        Commit a value to the DAG.
 612
 613        Parameters
 614        ----------
 615        value : Union[Node, Error, Any]
 616            Value to commit
 617        """
 618        # errors are committed as-is, everything else is a node
 619        if not isinstance(value, (Error, Node)):
 620            value = self.put(value)
 621        if isinstance(value, Node):
 622            value = value.ref
 623        self.ref = self.dml.runtime.commit(self._require_index_ref(), value, message=self.message, name=self.name)
 624        self.token = None  # Clear the working index since it's now committed
 625
 626    def freeze(self, message: str | None = None) -> "Dag":
 627        """Freeze this uncommitted DAG's runtime index for read-only inspection."""
 628        freeze_message = f"dag: {self.name}"
 629        if message:
 630            freeze_message = f"{freeze_message}\n{message}"
 631        self.token = self.dml.runtime.freeze(self._require_index_ref(), message=freeze_message)
 632        return self
 633
 634    def unfreeze(self) -> "Dag":
 635        """Unfreeze this uncommitted DAG's runtime index."""
 636        self.token = self.dml.runtime.unfreeze(self._require_index_ref())
 637        return self
 638
 639    def cancel(self, max_retries: int = 3):
 640        """Cancel the DAG's execution.
 641
 642        Parameters
 643        ----------
 644        max_retries : int, default=3
 645            Maximum cancellation retries after the initial attempt.
 646        """
 647        if self.token is None:
 648            raise DmlRepoError("Cannot cancel a committed DAG")
 649        logger.info(f"Cancelling execution {self.token}")
 650        self.dml.runtime.cancel(self.token, max_retries=max_retries)
 651        self.token = None  # Clear the index ref to indicate it's no longer active
 652        raise CancellationError("DAG execution cancelled")
 653
 654
 655class NodeError(Error):
 656    """A stored error enriched with the node that produced it."""
 657
 658    def __init__(self, error: Error, *, node_ref: Ref, dag: Dag):
 659        super().__init__(message=error.message, origin=error.origin, type=error.type, stack=list(error.stack))
 660        self.node_ref = node_ref
 661        self.dag = dag
 662
 663    def context(self) -> Dag:
 664        """Return the function DAG that recorded this node's failure."""
 665        dag = self.dag
 666        node_ref = self.node_ref
 667        while True:
 668            info = dag.dml.dag.describe_node(node_ref)
 669            if info["type"] == "FnNode":
 670                return Dag(dml=dag.dml, ref=info["dag"])
 671            if info["type"] != "ImportNode":
 672                raise DmlRepoError(f"Node {node_ref} does not resolve to a failed function call")
 673            dag = Dag(dml=dag.dml, ref=info["dag"])
 674            node_ref = info["node"]
 675
 676
 677@dataclass(frozen=True)
 678class Node:  # noqa: F811
 679    """
 680    Representation of a node in a DaggerML DAG.
 681
 682    Parameters
 683    ----------
 684    dag : Dag
 685        Parent DAG
 686    ref : Ref
 687        Node reference
 688    """
 689
 690    dag: Dag
 691    ref: Ref
 692    _info: dict = field(default_factory=dict)
 693
 694    def __repr__(self):
 695        ref_id = self.ref if isinstance(self.ref, Error) else self.ref.to
 696        return f"{self.__class__.__name__}({ref_id})"
 697
 698    def __hash__(self):
 699        return hash(self.ref)
 700
 701    def __eq__(self, other):
 702        if not isinstance(other, Node):
 703            return NotImplemented
 704        return self.ref == other.ref
 705
 706    def context(self, *, root: bool = True) -> Dag:
 707        """
 708        Resolve the provenance context (DAG) for this node.
 709
 710        This follows import/function provenance while treating builtin
 711        collection-construction and selection DAGs as transparent.
 712
 713        Parameters
 714        ----------
 715        root : bool, default=True
 716            If False, return the nearest sub-DAG in which this value exists as a
 717            proper node across a non-builtin import/function boundary. If True,
 718            continue recursively until provenance no longer crosses a
 719            non-builtin import/function boundary and return that first rooted
 720            context.
 721
 722        Returns
 723        -------
 724        Dag
 725            The nearest or rooted provenance DAG for this node.
 726
 727        Examples
 728        --------
 729        >>> source = new(dml=dml, name="source", message="source")
 730        >>> answer = source.put(42, name="answer")
 731        >>> payload = source.put({"answer": answer}, name="payload")
 732        >>> source.commit(payload)
 733        >>> consumer = new(dml=dml, name="consumer", message="consumer")
 734        >>> imported = consumer.require("source", "payload", name="payload")
 735        >>> consumer.commit(imported)
 736        >>> loaded = load("consumer", dml=dml)
 737        >>> loaded.result["answer"].context() == source
 738        True
 739        """
 740        return _resolve_context(self, (), root=root)
 741
 742    @property
 743    def type(self):
 744        """Get the data type of the node."""
 745        return self._info["data_type"]
 746
 747    @overload
 748    def value(self: "ScalarNode") -> Scalar: ...
 749    @overload
 750    def value(self: "ListNode") -> list: ...
 751    @overload
 752    def value(self: "DictNode") -> dict: ...
 753    @overload
 754    def value(self: "RunnableNode") -> Runnable: ...
 755    @overload
 756    def value(self: "Node") -> Any: ...
 757    def value(self):
 758        """
 759        Get the concrete value of this node.
 760
 761        Returns
 762        -------
 763        Any
 764            The actual value represented by this node
 765        """
 766        value = self.dag.dml.dag.get_node(self.ref, recursive=True)
 767        if isinstance(value, Error):
 768            raise NodeError(value, node_ref=self.ref, dag=self.dag)
 769        return value
 770
 771    def __call__(self, *args, name=None, sleep=None, timeout=-1, **kw) -> "Node":
 772        raise TypeError(f"Node of type '{self.type}' is not callable")
 773
 774
 775class ScalarNode(Node):
 776    pass
 777
 778
 779class RunnableNode(Node):
 780    def __call__(self, *args, name=None, sleep=None, timeout=-1) -> "Node":
 781        """
 782        Call this node as a function.
 783
 784        Parameters
 785        ----------
 786        *args : Any
 787            Arguments to pass to the function
 788        name : str, optional
 789            Name for the result node
 790        sleep : callable, optional
 791            A nullary function that returns sleep time in milliseconds
 792        timeout : int, default=-1
 793            Maximum time to wait in milliseconds. -1 means wait forever.
 794
 795        Returns
 796        -------
 797        Node
 798            Result node
 799
 800        Raises
 801        ------
 802        TimeoutError
 803            If the function call exceeds the timeout
 804        Error
 805            If the function returns an error
 806        """
 807        return self.dag.call(self, *args, name=name, sleep=sleep, timeout=timeout)
 808
 809
 810@dataclass(frozen=True)
 811class Projection:
 812    dag: Dag
 813    base: Node
 814    path: tuple[ProjectionStep, ...]
 815    _info: dict = field(default_factory=dict)
 816
 817    def __repr__(self):
 818        return f"Projection({self.base!r}, path={self.path!r})"
 819
 820    @classmethod
 821    def from_step(cls, base: Node, step: ProjectionStep) -> "Projection":
 822        return cls(dag=base.dag, base=base, path=(step,))
 823
 824    def _extend(self, step: ProjectionStep) -> "Projection":
 825        return Projection(dag=self.dag, base=self.base, path=(*self.path, step))
 826
 827    def value(self):
 828        return _apply_projection_path(self.base.value(), self.path)
 829
 830    def context(self, *, root: bool = True) -> Dag:
 831        return _resolve_context(self.base, self.path, root=root)
 832
 833    def __call__(self, *args, **kwargs):
 834        raise TypeError(f"Projection of type '{self.type}' is not callable")
 835
 836    @property
 837    def type(self):
 838        if "data_type" in self._info:
 839            return self._info["data_type"]
 840        return _info_for_value(self.value())["data_type"]
 841
 842    def __len__(self):
 843        if "length" in self._info:
 844            return self._info["length"]
 845        value = self.value()
 846        if not isinstance(value, (list, dict)):
 847            raise TypeError(f"Object of type '{type(value).__name__}' has no len()")
 848        return len(value)
 849
 850    def __iter__(self):
 851        value = self.value()
 852        if isinstance(value, list):
 853            for i in range(len(value)):
 854                yield self[i]
 855            return
 856        if isinstance(value, dict):
 857            yield from self.keys()
 858            return
 859        raise TypeError(f"Object of type '{type(value).__name__}' is not iterable")
 860
 861    def __getitem__(self, key: ProjectionStep | slice) -> "Projection":
 862        value = self.value()
 863        if isinstance(value, dict):
 864            if not isinstance(key, str):
 865                raise TypeError(f"Dict keys must be strings but got {type(key).__name__}")
 866            step = cast(ProjectionStep, key)
 867        elif isinstance(value, list):
 868            if not isinstance(key, (int, slice)):
 869                raise TypeError(f"List indices must be integers or slices but got {type(key).__name__}")
 870            step = _normalize_projection_step(cast(ProjectionStep | slice, key), length=len(value))
 871        else:
 872            raise TypeError(f"Cannot project into object of type '{type(value).__name__}'")
 873        projected = self._extend(step)
 874        return Projection(dag=self.dag, base=self.base, path=projected.path, _info=_info_for_value(projected.value()))
 875
 876    def keys(self) -> list[str]:
 877        value = self.value()
 878        if not isinstance(value, dict):
 879            raise TypeError(f"Cannot get keys of type: {type(value).__name__}")
 880        return sorted(value.keys())
 881
 882
 883class CollectionNode(Node):  # noqa: F811
 884    """
 885    Representation of a collection node in a DaggerML DAG.
 886
 887    Parameters
 888    ----------
 889    dag : Dag
 890        Parent DAG
 891    ref : Ref
 892        Node reference
 893    """
 894
 895    def contains(self, item, *, name=None) -> "ScalarNode":
 896        """
 897        For collection nodes, checks to see if `item` is in `self`
 898
 899        Returns
 900        -------
 901        Node
 902            Node with the boolean of is `item` in `self`
 903        """
 904        item_ref = item.ref if isinstance(item, Node) else item
 905        result = self.dag._call_builtin("daggerml:contains", self.ref, item_ref, name=name)
 906        return cast(ScalarNode, _make_node(self.dag, result))
 907
 908    def __contains__(self, item):
 909        return self.contains(item).value()  # has to return boolean
 910
 911    def __len__(self):  # python requires this to be an int
 912        """
 913        Get the node's length
 914
 915        Returns
 916        -------
 917        Node
 918            Node with the length of the collection
 919
 920        Raises
 921        ------
 922        Error
 923            If the node isn't a collection (e.g. list or dict).
 924        """
 925        return self._info["length"]
 926
 927
 928class ListNode(CollectionNode):  # noqa: F811
 929    """
 930    Representation of a collection node in a DaggerML DAG.
 931
 932    Parameters
 933    ----------
 934    dag : Dag
 935        Parent DAG
 936    ref : Ref
 937        Node reference
 938    """
 939
 940    @overload
 941    def __getitem__(self, key: Union[slice, list[int]]) -> Union["ListNode", Projection]: ...
 942    @overload
 943    def __getitem__(self, key: Union[int, "Node"]) -> Union["Node", Projection]: ...
 944    def __getitem__(self, key: Union[slice, list[int], int, "Node"]) -> Union["Node", Projection]:
 945        if self.dag.ref is not None:
 946            if isinstance(key, Node):
 947                raise TypeError("Committed list projections require concrete int or slice keys")
 948            step = _normalize_projection_step(key, length=len(self))
 949            return Projection(
 950                dag=self.dag,
 951                base=self,
 952                path=(step,),
 953                _info=_info_for_value(_apply_projection_step(self.value(), step)),
 954            )
 955        if isinstance(key, slice):
 956            if key.step is not None:
 957                raise ValueError("Slice step is not supported")
 958            start = key.start if key.start is not None else 0
 959            stop = key.stop if key.stop is not None else len(self)
 960            key = [start, stop]
 961        return _make_node(self.dag, self.dag._call_builtin("daggerml:get", self.ref, key))
 962
 963    def __iter__(self):
 964        """
 965        Iterate over the node's values (items if it's a list, and keys if it's a
 966        dict)
 967
 968        Returns
 969        -------
 970        Node
 971            Result node
 972
 973        Raises
 974        ------
 975        Error
 976            If the node isn't a collection (e.g. list or dict).
 977        """
 978        for i in range(len(self)):
 979            yield self[i]
 980
 981    def conj(self, item, *, name=None) -> "ListNode":
 982        """
 983        For a list node, append an item
 984
 985        Returns
 986        -------
 987        Node
 988            Node containing the new collection
 989
 990        Notes
 991        -----
 992        `append` is an alias `conj`
 993        """
 994        item_ref = item.ref if isinstance(item, Node) else item
 995        resp = self.dag._call_builtin("daggerml:conj", self.ref, item_ref, name=name)
 996        return cast(ListNode, _make_node(self.dag, resp))
 997
 998    def append(self, item, *, name=None) -> "ListNode":
 999        """
1000        For a list node, append an item
1001
1002        Returns
1003        -------
1004        Node
1005            Node containing the new collection
1006
1007        See Also
1008        --------
1009        conj : The main implementation
1010        """
1011        return self.conj(item, name=name)
1012
1013
1014class DictNode(CollectionNode):  # noqa: F811
1015    def __getitem__(self, key: Union[str, "Node"]) -> Union["Node", Projection]:
1016        if self.dag.ref is not None:
1017            if not isinstance(key, str):
1018                raise TypeError(f"Dict keys must be strings but got {type(key).__name__}")
1019            return Projection(dag=self.dag, base=self, path=(key,), _info=_info_for_value(self.value()[key]))
1020        return _make_node(self.dag, self.dag._call_builtin("daggerml:get", self.ref, key))
1021
1022    def keys(self) -> list[str]:
1023        """
1024        Get the keys of a dictionary node.
1025
1026        Parameters
1027        ----------
1028        name : str, optional
1029            Name for the result node
1030
1031        Returns
1032        -------
1033        list[str]
1034            List of keys in the dictionary node
1035        """
1036        return self._info["keys"].copy()
1037
1038    def __iter__(self):
1039        """
1040        Iterate over the node's values (items if it's a list, and keys if it's a
1041        dict)
1042
1043        Returns
1044        -------
1045        Node
1046            Result node
1047
1048        Raises
1049        ------
1050        Error
1051            If the node isn't a collection (e.g. list or dict).
1052        """
1053        for k in self.keys():
1054            yield k
1055
1056    def get(self, key, default=None, *, name=None) -> "Node":
1057        """
1058        For a dict node, return the value for key if key exists, else default.
1059
1060        If default is not given, it defaults to None, so that this method never raises a KeyError.
1061        """
1062        return _make_node(self.dag, self.dag._call_builtin("daggerml:get", self.ref, key, default, name=name))
1063
1064    def items(self) -> Iterator[tuple[str, "Node|Projection"]]:
1065        """
1066        Iterate over key-value pairs of a dictionary node.
1067
1068        Returns
1069        -------
1070        Iterator[tuple[Node, Node]]
1071            Iterator over (key, value) pairs
1072        """
1073        if self.type != "dict":
1074            raise Error(f"Cannot iterate items of type: {self.type}", origin="dml", type="TypeError")
1075        for k in self:
1076            yield k, self[k]
1077
1078    def values(self) -> list["Node|Projection"]:
1079        """
1080        Get the values of a dictionary node.
1081
1082        Parameters
1083        ----------
1084        name : str, optional
1085            Name for the result node
1086
1087        Returns
1088        -------
1089        list[Node]
1090            List of values in the dictionary node
1091        """
1092        return [self[k] for k in self]
1093
1094    def assoc(self, key, value, *, name=None) -> "DictNode":
1095        """
1096        For a dict node, associate a new value into the map
1097
1098        Returns
1099        -------
1100        Node
1101            Node containing the new dict
1102        """
1103        value_ref = value.ref if isinstance(value, Node) else value
1104        resp = self.dag._call_builtin("daggerml:assoc", self.ref, key, value_ref, name=name)
1105        return cast(DictNode, _make_node(self.dag, resp))
1106
1107    def update(self, update) -> "DictNode":
1108        """
1109        For a dict node, update like python dicts
1110
1111        Returns
1112        -------
1113        Node
1114            Node containing the new collection
1115
1116        Notes
1117        -----
1118        calls `assoc` iteratively for k, v pairs in update.
1119
1120        See Also
1121        --------
1122        assoc : The main implementation
1123        """
1124        for k, v in update.items():
1125            self = self.assoc(k, v)
1126        return self
1127
1128
1129################################################################################
1130################## Codec system for encoding literals in DAGs ##################
1131################################################################################
1132
1133LITERAL_CODEC_ENTRYPOINT_GROUP = "daggerml.codecs"
1134_codecs: list[tuple[int, int, "LiteralCodec"]] = []
1135_plugins_loaded = False
1136_lock = RLock()
1137
1138
1139class LiteralCodec(Protocol):
1140    def can_encode(self, value: Any) -> bool: ...
1141
1142    def encode(self, value: Any, dag: "Dag") -> Any: ...
1143
1144
1145class CodecError(Error):
1146    def __init__(self, message: str):
1147        super().__init__(message, origin="dml-codec", type="codec-error")
1148
1149
1150def _entry_points(group=LITERAL_CODEC_ENTRYPOINT_GROUP) -> list[metadata.EntryPoint]:
1151    points = metadata.entry_points()
1152    result = list(points.select(group=group))
1153    result.sort(key=lambda ep: (ep.name, ep.value))
1154    return result
1155
1156
1157def ensure_literal_codec_plugins_loaded() -> None:
1158    global _plugins_loaded
1159    if _plugins_loaded:
1160        return
1161    with _lock:
1162        codec_seq = 0
1163        if _plugins_loaded:
1164            return
1165        loaded = []
1166        for entry_point in _entry_points():
1167            try:
1168                registrations = entry_point.load()()
1169                for item in registrations:
1170                    priority, codec = item
1171                    codec_seq += 1
1172                    loaded.append((priority, codec_seq, codec))
1173            except Exception as e:
1174                msg = f"Literal codec plugin '{entry_point.name} ({entry_point.value})' failed: {e}"
1175                raise CodecError(msg) from None
1176        loaded.sort(key=lambda item: (-item[0], item[1]))
1177        _codecs.extend(loaded)
1178        _plugins_loaded = True
1179
1180
1181def iter_codecs() -> Iterator[LiteralCodec]:
1182    ensure_literal_codec_plugins_loaded()
1183    yield from [codec for _priority, _seq, codec in _codecs]
1184
1185
1186def apply_codec(value: Any, *, dag: Dag) -> Any:
1187    for codec in iter_codecs():
1188        try:
1189            if codec.can_encode(value):
1190                resp = codec.encode(value, dag)
1191                if isinstance(resp, type(value)):
1192                    codec_name = codec.__class__.__name__
1193                    msg = f"Literal codec {codec_name} encoded {value.__class__.__name__} to {resp.__class__.__name__}."
1194                    raise CodecError(msg)
1195                return resp
1196        except Exception as e:
1197            if isinstance(e, DmlRepoError):
1198                raise
1199            raise CodecError(f"Literal codec {codec.__class__.__name__} failed: {e}") from e
1200    raise CodecError(f"No codec found for value of type {type(value).__name__}")
1201
1202
1203def apply_codecs(value: Any, *, dag: Dag) -> Any:
1204    while not isinstance(value, (*get_args(Scalar), *get_args(Collection), Ref)):
1205        value = apply_codec(value, dag=dag)
1206    if isinstance(value, list):
1207        return [apply_codecs(v, dag=dag) for v in value]
1208    if isinstance(value, dict):
1209        return {k: apply_codecs(v, dag=dag) for k, v in value.items()}
1210    if isinstance(value, Uri):
1211        return Uri(apply_codecs(value.uri, dag=dag))
1212    if isinstance(value, Runnable):
1213        target = apply_codecs(value.target, dag=dag)
1214        sub = apply_codecs(value.sub, dag=dag)
1215        kwargs = {k: apply_codecs(v, dag=dag) for k, v in value.kwargs.items()}
1216        return Runnable(target=target, adapter=value.adapter, kwargs=kwargs, sub=sub)
1217    return value
1218
1219
1220class MiscPyTypeCodec:
1221    def can_encode(self, value: Any) -> bool:
1222        return isinstance(value, Mapping) or (
1223            isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray))
1224        )
1225
1226    def encode(self, value: Sequence | Mapping, dag: Dag) -> Any:
1227        if isinstance(value, Mapping):
1228            return {k: apply_codecs(v, dag=dag) for k, v in value.items()}
1229        if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
1230            return [apply_codecs(v, dag=dag) for v in value]
1231
1232
1233class NodeCodec:
1234    def can_encode(self, value: Any) -> bool:
1235        return isinstance(value, Node)
1236
1237    def encode(self, value: "Node", dag: Dag) -> Ref:
1238        assert dag.token is not None, "DAG must have a token to encode nodes"
1239        if value.dag.token is not None and value.dag.token == dag.token:
1240            return value.ref
1241        if value.dag.ref is None:
1242            raise CodecError("Cannot encode node from uncommitted DAG in a different index")
1243        try:
1244            return dag.dml.runtime.put_import(dag._require_index_ref(), value.dag.ref, node=value.ref, name=None)
1245        except Exception as e:
1246            raise CodecError(f"Failed to encode cross-dag node import: {e}") from e
1247
1248
1249class ProjectionCodec:
1250    def can_encode(self, value: Any) -> bool:
1251        return isinstance(value, Projection)
1252
1253    def encode(self, value: "Projection", dag: Dag) -> Ref:
1254        assert dag.token is not None, "DAG must have a token to encode projections"
1255        node_ref = NodeCodec().encode(value.base, dag)
1256        for step in value.path:
1257            node_ref = dag._call_builtin("daggerml:get", node_ref, step)
1258        return node_ref
1259
1260
1261def codecs() -> list:
1262    return [(0, NodeCodec()), (0, MiscPyTypeCodec()), (0, ProjectionCodec())]