data-annotations 2.1.2__py3-none-any.whl

data_annotations/provenance/runtime.py

@@ -0,0 +1,248 @@
+import json
+import os
+import socket
+import subprocess
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable
+from urllib.parse import urlencode
+from urllib.request import Request, urlopen
+
+
+def utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def _connection_file_from_argv(argv: list[str]) -> str | None:
+    for index, arg in enumerate(argv[1:], start=1):
+        if arg.startswith("--f="):
+            return arg.split("=", 1)[1]
+        if arg == "-f" and index + 1 < len(argv):
+            return argv[index + 1]
+    try:
+        from ipykernel.connect import get_connection_file
+
+        return get_connection_file()
+    except Exception:
+        return None
+
+
+def _kernel_id_from_connection_file(connection_file: str | Path | None) -> str | None:
+    if connection_file is None:
+        return None
+    stem = Path(connection_file).stem
+    if not stem.startswith("kernel-"):
+        return None
+    return stem.removeprefix("kernel-")
+
+
+def _iter_jupyter_server_info_files(runtime_dir: Path) -> list[Path]:
+    return sorted(
+        [
+            *runtime_dir.glob("jpserver-*.json"),
+            *runtime_dir.glob("nbserver-*.json"),
+        ]
+    )
+
+
+def _current_ipython():
+    try:
+        from IPython import get_ipython
+    except Exception:
+        return None
+    return get_ipython()
+
+
+def _normalize_session_path(
+    path: str | Path | None,
+    *,
+    cwd: Path | None = None,
+) -> Path | None:
+    if path is None:
+        return None
+    path = Path(path).expanduser()
+    if path.is_absolute():
+        return path.resolve()
+    base_dir = Path.cwd() if cwd is None else cwd
+    return (base_dir / path).resolve()
+
+
+def _notebook_path_from_session_name(
+    *,
+    env: dict[str, str] | os._Environ[str] | None = None,
+    current_ipython_fn: Callable[[], Any] | None = None,
+    cwd: Path | None = None,
+) -> Path | None:
+    env = os.environ if env is None else env
+    current_ipython_fn = (
+        _current_ipython if current_ipython_fn is None else current_ipython_fn
+    )
+    session_name = env.get("JPY_SESSION_NAME")
+    if session_name:
+        return _normalize_session_path(session_name, cwd=cwd)
+
+    shell = current_ipython_fn()
+    if shell is None:
+        return None
+
+    session_name = shell.user_ns.get("__session__")
+    if not isinstance(session_name, str) or not session_name:
+        return None
+
+    return _normalize_session_path(session_name, cwd=cwd)
+
+
+def _resolve_notebook_path(
+    argv: list[str],
+    *,
+    env: dict[str, str] | os._Environ[str] | None = None,
+    current_ipython_fn: Callable[[], Any] | None = None,
+    request_cls=None,
+    urlopen_fn=None,
+    cwd: Path | None = None,
+) -> Path | None:
+    env = os.environ if env is None else env
+    current_ipython_fn = (
+        _current_ipython if current_ipython_fn is None else current_ipython_fn
+    )
+    request_cls = Request if request_cls is None else request_cls
+    urlopen_fn = urlopen if urlopen_fn is None else urlopen_fn
+    if not argv or Path(argv[0]).name != "ipykernel_launcher.py":
+        return None
+
+    notebook_path = _notebook_path_from_session_name(
+        env=env,
+        current_ipython_fn=current_ipython_fn,
+        cwd=cwd,
+    )
+    if notebook_path is not None:
+        return notebook_path
+
+    connection_file = _connection_file_from_argv(argv)
+    kernel_id = _kernel_id_from_connection_file(connection_file)
+    if connection_file is None or kernel_id is None:
+        return None
+
+    runtime_dir = Path(connection_file).expanduser().resolve().parent
+    for server_info_path in _iter_jupyter_server_info_files(runtime_dir):
+        try:
+            server_info = json.loads(server_info_path.read_text(encoding="utf-8"))
+        except Exception:
+            continue
+
+        server_url = server_info.get("url")
+        if not server_url:
+            continue
+
+        base_url = server_info.get("base_url", "/").rstrip("/")
+        sessions_url = f"{server_url.rstrip('/')}{base_url}/api/sessions"
+        token = server_info.get("token")
+        if token:
+            sessions_url = f"{sessions_url}?{urlencode({'token': token})}"
+
+        try:
+            request = request_cls(sessions_url, headers={"Accept": "application/json"})
+            with urlopen_fn(request, timeout=0.2) as response:
+                sessions = json.loads(response.read().decode("utf-8"))
+        except Exception:
+            continue
+
+        for session in sessions:
+            session_kernel_id = (session.get("kernel") or {}).get("id")
+            if session_kernel_id != kernel_id:
+                continue
+
+            notebook_path = (
+                (session.get("notebook") or {}).get("path")
+            ) or session.get("path")
+            if not notebook_path:
+                return None
+
+            root_dir = server_info.get("root_dir") or server_info.get("notebook_dir")
+            if root_dir:
+                return (Path(root_dir) / notebook_path).resolve()
+            return Path(notebook_path).expanduser().resolve()
+
+    return None
+
+
+def _git_repo_root(
+    path: str | Path,
+    *,
+    subprocess_module=subprocess,
+) -> Path | None:
+    candidate = Path(path).expanduser().resolve()
+    cwd = candidate if candidate.is_dir() else candidate.parent
+    try:
+        repo_root = subprocess_module.check_output(
+            ["git", "-C", str(cwd), "rev-parse", "--show-toplevel"],
+            text=True,
+            stderr=subprocess_module.DEVNULL,
+        ).strip()
+    except Exception:
+        return None
+
+    if not repo_root:
+        return None
+    return Path(repo_root).resolve()
+
+
+def infer_script_repo_path(script: str | Path | None) -> str | None:
+    if script is None:
+        return None
+    script_path = Path(script).expanduser().resolve()
+    repo_root = _git_repo_root(script_path)
+    if repo_root is None:
+        return None
+    try:
+        return str(script_path.relative_to(repo_root))
+    except ValueError:
+        return None
+
+
+def capture_runtime_info(
+    *,
+    argv: list[str] | None = None,
+    utcnow_fn: Callable[[], datetime] | None = None,
+    gethostname: Callable[[], str] | None = None,
+    env: dict[str, str] | os._Environ[str] | None = None,
+    current_ipython_fn: Callable[[], Any] | None = None,
+    request_cls=None,
+    urlopen_fn=None,
+    cwd: Path | None = None,
+) -> dict[str, Any]:
+    utcnow_fn = utcnow if utcnow_fn is None else utcnow_fn
+    gethostname = socket.gethostname if gethostname is None else gethostname
+    env = os.environ if env is None else env
+    current_ipython_fn = (
+        _current_ipython if current_ipython_fn is None else current_ipython_fn
+    )
+    request_cls = Request if request_cls is None else request_cls
+    urlopen_fn = urlopen if urlopen_fn is None else urlopen_fn
+    argv = list(sys.argv) if argv is None else list(argv)
+    notebook_path = _resolve_notebook_path(
+        argv,
+        env=env,
+        current_ipython_fn=current_ipython_fn,
+        request_cls=request_cls,
+        urlopen_fn=urlopen_fn,
+        cwd=cwd,
+    )
+    if notebook_path is not None:
+        script = str(notebook_path)
+        command = [str(notebook_path)]
+    else:
+        script = os.path.abspath(argv[0]) if argv else None
+        command = argv
+    script_repo_path = infer_script_repo_path(script)
+
+    return {
+        "created_at": utcnow_fn(),
+        "hostname": gethostname(),
+        "username": env.get("USER", "unknown"),
+        "script": script,
+        "script_repo_path": script_repo_path,
+        "command": command,
+        "slurm_job_id": env.get("SLURM_JOB_ID"),
+    }

data_annotations/provenance/writers.py

@@ -0,0 +1,206 @@
+import hashlib
+from pathlib import Path
+from typing import Any, Callable
+from urllib.parse import urlparse
+
+from . import git, runtime
+from .models import (
+    ArtifactKind,
+    BaseProvenance,
+    DirectoryManifest,
+    FileManifest,
+    ProducedFile,
+)
+
+
+def sha256_file(path: str | Path, chunk_size: int = 1024 * 1024) -> str:
+    path = Path(path)
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(chunk_size), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def callable_name(fn: Callable[..., Any] | None) -> str | None:
+    if fn is None:
+        return None
+    module = getattr(fn, "__module__", fn.__class__.__module__)
+    qualname = getattr(fn, "__qualname__", fn.__class__.__qualname__)
+    return f"{module}.{qualname}"
+
+
+def _normalize_local_path(path: str | Path) -> str:
+    return str(Path(path).expanduser().resolve())
+
+
+def _is_uri_like(value: str) -> bool:
+    parsed = urlparse(value)
+    return bool(parsed.scheme and (parsed.netloc or "://" in value))
+
+
+def _normalize_input(value: str | Path) -> str:
+    if isinstance(value, str) and _is_uri_like(value):
+        return value
+    return _normalize_local_path(value)
+
+
+def _normalize_produced_file(produced_file: ProducedFile) -> ProducedFile:
+    return produced_file.model_copy(
+        update={"path": _normalize_local_path(produced_file.path)}
+    )
+
+
+def _manifest_context(
+    *,
+    params: dict[str, Any] | None = None,
+    inputs: list[str] | None = None,
+    function: Callable[..., Any] | None = None,
+    capture_mode: str = "runtime",
+    overrides: dict[str, Any] | None = None,
+    normalize_inputs: bool = True,
+) -> dict[str, Any]:
+    context: dict[str, Any] = {
+        **runtime.capture_runtime_info(),
+        **git.capture_git_info(),
+        "capture_mode": capture_mode,
+        "params": params or {},
+        "inputs": (
+            [_normalize_input(value) for value in (inputs or [])]
+            if normalize_inputs
+            else list(inputs or [])
+        ),
+    }
+    if function is not None:
+        context["function"] = callable_name(function)
+    if overrides:
+        context.update(overrides)
+    return context
+
+
+def _build_directory_manifest(
+    output_dir: str | Path,
+    *,
+    produced_files: list[ProducedFile],
+    params: dict[str, Any] | None = None,
+    inputs: list[str] | None = None,
+    function: Callable[..., Any] | None = None,
+    capture_mode: str = "runtime",
+    overrides: dict[str, Any] | None = None,
+    normalize_inputs: bool = True,
+) -> DirectoryManifest:
+    normalized_output_dir = Path(_normalize_local_path(output_dir))
+    return DirectoryManifest(
+        **_manifest_context(
+            params=params,
+            inputs=inputs,
+            function=function,
+            capture_mode=capture_mode,
+            overrides=overrides,
+            normalize_inputs=normalize_inputs,
+        ),
+        output_dir=str(normalized_output_dir),
+        produced_files=[_normalize_produced_file(item) for item in produced_files],
+    )
+
+
+def _build_file_manifest(
+    artifact_path: str | Path,
+    *,
+    artifact_kind: ArtifactKind = "other",
+    params: dict[str, Any] | None = None,
+    inputs: list[str] | None = None,
+    function: Callable[..., Any] | None = None,
+    capture_mode: str = "runtime",
+    overrides: dict[str, Any] | None = None,
+    normalize_inputs: bool = True,
+) -> FileManifest:
+    normalized_artifact_path = Path(_normalize_local_path(artifact_path))
+    return FileManifest(
+        **_manifest_context(
+            params=params,
+            inputs=inputs,
+            function=function,
+            capture_mode=capture_mode,
+            overrides=overrides,
+            normalize_inputs=normalize_inputs,
+        ),
+        artifact_path=str(normalized_artifact_path),
+        artifact_kind=artifact_kind,
+        artifact_sha256=(
+            sha256_file(normalized_artifact_path)
+            if normalized_artifact_path.exists()
+            else None
+        ),
+    )
+
+
+def _write_manifest(manifest: FileManifest | DirectoryManifest, path: Path) -> Path:
+    path.write_text(manifest.model_dump_json(indent=2), encoding="utf-8")
+    return path
+
+
+def write_directory_manifest(
+    output_dir: str | Path,
+    *,
+    produced_files: list[ProducedFile],
+    params: dict[str, Any] | None = None,
+    inputs: list[str] | None = None,
+    function: Callable[..., Any] | None = None,
+    filename: str = "manifest.json",
+) -> Path:
+    from data_annotations.annotations import (
+        DirectoryAnnotationDocument,
+        DirectoryArtifactSubject,
+    )
+    from data_annotations.description import ArtifactDescription, DirectoryDescription
+
+    normalized_output_dir = Path(_normalize_local_path(output_dir))
+    normalized_output_dir.mkdir(parents=True, exist_ok=True)
+
+    manifest = _build_directory_manifest(
+        normalized_output_dir,
+        produced_files=produced_files,
+        params=params,
+        inputs=inputs,
+        function=function,
+    )
+    document = DirectoryAnnotationDocument(
+        subject=DirectoryArtifactSubject(
+            path=manifest.output_dir,
+            produced_files=list(manifest.produced_files),
+        ),
+        provenance=BaseProvenance.model_validate(
+            manifest.model_dump(exclude={"output_dir", "produced_files"})
+        ),
+        description=DirectoryDescription(
+            artifacts=[
+                ArtifactDescription(path=item.path) for item in manifest.produced_files
+            ],
+            description_updated_at=manifest.created_at,
+        ),
+    )
+    path = normalized_output_dir / filename
+    path.write_text(document.model_dump_json(indent=2), encoding="utf-8")
+    return path
+
+
+def write_file_manifest(
+    artifact_path: str | Path,
+    *,
+    artifact_kind: ArtifactKind = "other",
+    params: dict[str, Any] | None = None,
+    inputs: list[str] | None = None,
+    function: Callable[..., Any] | None = None,
+    suffix: str = ".meta.json",
+) -> Path:
+    from data_annotations.annotations import write_file_annotation
+
+    return write_file_annotation(
+        artifact_path,
+        artifact_kind=artifact_kind,
+        params=params,
+        inputs=inputs,
+        function=function,
+        suffix=suffix,
+    )