ygg 0.1.56__py3-none-any.whl → 0.1.57__py3-none-any.whl

{ygg-0.1.56.dist-info → ygg-0.1.57.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.56
+Version: 0.1.57
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 License:                                  Apache License

{ygg-0.1.56.dist-info → ygg-0.1.57.dist-info}/RECORD

@@ -1,14 +1,12 @@
-ygg-0.1.56.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+ygg-0.1.57.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
 yggdrasil/__init__.py,sha256=4-ghPak2S6zfMqmnlxW2GCgPb5s79znpKa2hGEGXcE4,24
-yggdrasil/exceptions.py,sha256=NEpbDFn-8ZRsLiEgJicCwrTHNMWAGtdrTJzosfAeVJo,82
-yggdrasil/version.py,sha256=c0ITmemMU7anWgDBUD3t_BAhA3li2gt3XgswCbHv1oU,22
+yggdrasil/version.py,sha256=mM67BdyYZ17u9xAi4WRzFQM2e6yfmX4MPd36R3L920M,22
 yggdrasil/databricks/__init__.py,sha256=skctY2c8W-hI81upx9F_PWRe5ishL3hrdiTuizgDjdw,152
-yggdrasil/databricks/ai/__init__.py,sha256=Mkp70UOVBzDQvdPNsqncHcyzxe5PnSGYE_bHnYxA1eA,21
-yggdrasil/databricks/ai/loki.py,sha256=HyVWxzJgfW03YO6TMOTJ1oNvrBJovnqYDn_MeNV6Ni0,11989
+yggdrasil/databricks/ai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+yggdrasil/databricks/ai/loki.py,sha256=1AhMOquMTsZZGYw5cGoXn-QQhdBRMXM9ZRPEUAv4Y3k,1216
 yggdrasil/databricks/compute/__init__.py,sha256=NvdzmaJSNYY1uJthv1hHdBuNu3bD_-Z65DWnaJt9yXg,289
 yggdrasil/databricks/compute/cluster.py,sha256=YomLfvB0oxbgl6WDgBRxI1UXsxwlEbR6gq3FUbPHscY,44199
-yggdrasil/databricks/compute/exceptions.py,sha256=Ug0ioxu5m2atdTX4OLH0s4R4dylHNxEdn7VhQI66b5M,209
-yggdrasil/databricks/compute/execution_context.py,sha256=mhcwSvKTxgcUHdb7huSEjCVU_feiXSGq0JLyLXldjQM,23952
+yggdrasil/databricks/compute/execution_context.py,sha256=jIV6uru2NeX3O5lg-3KEqmXtLxxq45CFgkBQgQIIOHQ,23327
 yggdrasil/databricks/compute/remote.py,sha256=yicEhyQypssRa2ByscO36s3cBkEgORFsRME9aaq91Pc,3045
 yggdrasil/databricks/jobs/__init__.py,sha256=snxGSJb0M5I39v0y3IR-uEeSlZR248cQ_4DJ1sYs-h8,154
 yggdrasil/databricks/jobs/config.py,sha256=9LGeHD04hbfy0xt8_6oobC4moKJh4_DTjZiK4Q2Tqjk,11557
@@ -20,11 +18,11 @@ yggdrasil/databricks/sql/types.py,sha256=5G-BM9_eOsRKEMzeDTWUsWW5g4Idvs-czVCpOCr
 yggdrasil/databricks/sql/warehouse.py,sha256=1J0dyQLJb-OS1_1xU1eAVZ4CoL2-FhFeowKSvU3RzFc,9773
 yggdrasil/databricks/workspaces/__init__.py,sha256=dv2zotoFVhNFlTCdRq6gwf5bEzeZkOZszoNZMs0k59g,114
 yggdrasil/databricks/workspaces/filesytem.py,sha256=Z8JXU7_XUEbw9fpTQT1avRQKi-IAP2KemXBMPkUoY4w,9805
-yggdrasil/databricks/workspaces/io.py,sha256=PAoxIxYvTC162Dx2qL2hk8oAdt8BnYrQ3jJHcJm4VkA,33116
-yggdrasil/databricks/workspaces/path.py,sha256=h1j3bvjwKcDhJvlU_kAaLcLVz4jrdaWgjqPQMarZRHU,55233
+yggdrasil/databricks/workspaces/io.py,sha256=hErGeSKJ9XpSUvlYAAckh_8IKQQmGeDOqbdl2rh9Fbs,33240
+yggdrasil/databricks/workspaces/path.py,sha256=KkvLFHrps3UFr4ogYdESbJHEMfQBcWfWfXjlrv_7rTU,55180
 yggdrasil/databricks/workspaces/path_kind.py,sha256=rhWe1ky7uPD0du0bZSv2S4fK4C5zWd7zAF3UeS2iiPU,283
 yggdrasil/databricks/workspaces/volumes_path.py,sha256=s8CA33cG3jpMVJy5MILLlkEBcFg_qInDCF2jozLj1Fg,2431
-yggdrasil/databricks/workspaces/workspace.py,sha256=5DCPz5io_rmrpGNi5I6RChmyZ8kjlNUFGQl8mzQJThg,25511
+yggdrasil/databricks/workspaces/workspace.py,sha256=Tl1pYzTGNpjsPmHCJ62HoJvdzHGiZb43vQxrI3Sk7js,25233
 yggdrasil/dataclasses/__init__.py,sha256=_RkhfF3KC1eSORby1dzvBXQ0-UGG3u6wyUQWX2jq1Pc,108
 yggdrasil/dataclasses/dataclass.py,sha256=LxrCjwvmBnb8yRI_N-c31RHHxB4XoJPixmKg9iBIuaI,1148
 yggdrasil/libs/__init__.py,sha256=zdC9OU0Xy36CLY9mg2drxN6S7isPR8aTLzJA6xVIeLE,91
@@ -35,14 +33,14 @@ yggdrasil/libs/sparklib.py,sha256=FQ3W1iz2EIpQreorOiQuFt15rdhq2QhGEAWp8Zrbl9A,10
 yggdrasil/libs/extensions/__init__.py,sha256=mcXW5Li3Cbprbs4Ci-b5A0Ju0wmLcfvEiFusTx6xNjU,117
 yggdrasil/libs/extensions/polars_extensions.py,sha256=RTkGi8llhPJjX7x9egix7-yXWo2X24zIAPSKXV37SSA,12397
 yggdrasil/libs/extensions/spark_extensions.py,sha256=E64n-3SFTDgMuXwWitX6vOYP9ln2lpGKb0htoBLEZgc,16745
-yggdrasil/pyutils/__init__.py,sha256=a7LKg9xWZGtv2o5jr_7wrSZJ9ZmwolKWc--8iRgKlL4,225
+yggdrasil/pyutils/__init__.py,sha256=tl-LapAc71TV7RMgf2ftKwrzr8iiLOGHeJgA3RvO93w,293
 yggdrasil/pyutils/callable_serde.py,sha256=1XckmFO-ThP0MedxgXwB71u9jWUuhM1btOzW9gJ8w9g,23117
 yggdrasil/pyutils/equality.py,sha256=Xyf8D1dLUCm3spDEir8Zyj7O4US_fBJwEylJCfJ9slI,3080
 yggdrasil/pyutils/exceptions.py,sha256=ssKNm-rjhavHUOZmGA7_1Gq9tSHDrb2EFI-cnBuWgng,3388
 yggdrasil/pyutils/expiring_dict.py,sha256=pr2u25LGwPVbLfsLptiHGovUtYRRo0AMjaJtCtJl7nQ,8477
 yggdrasil/pyutils/modules.py,sha256=B7IP99YqUMW6-DIESFzBx8-09V1d0a8qrIJUDFhhL2g,11424
 yggdrasil/pyutils/parallel.py,sha256=ubuq2m9dJzWYUyKCga4Y_9bpaeMYUrleYxdp49CHr44,6781
-yggdrasil/pyutils/python_env.py,sha256=nmXXcIUSSbpKWHBnYfCY1QIW-LQ8IILV9Ceijufu5Pg,51086
+yggdrasil/pyutils/python_env.py,sha256=Gh5geFK9ABpyWEfyegGUfIJUoPxKwcH0pqLBiMrW9Rw,51103
 yggdrasil/pyutils/retry.py,sha256=gXBtn1DdmIYIUmGKOUr8-SUT7MOu97LykN2YR4uocgc,11917
 yggdrasil/requests/__init__.py,sha256=dMesyzq97_DmI765x0TwaDPEfsxFtgGNgchk8LvEN-o,103
 yggdrasil/requests/msal.py,sha256=s2GCyzbgFdgdlJ1JqMrZ4qYVbmoG46-ZOTcaVQhZ-sQ,9220
@@ -61,8 +59,8 @@ yggdrasil/types/cast/registry.py,sha256=OOqIfbIjPH-a3figvu-zTvEtUDTEWhe2xIl3cCA4
 yggdrasil/types/cast/spark_cast.py,sha256=_KAsl1DqmKMSfWxqhVE7gosjYdgiL1C5bDQv6eP3HtA,24926
 yggdrasil/types/cast/spark_pandas_cast.py,sha256=BuTiWrdCANZCdD_p2MAytqm74eq-rdRXd-LGojBRrfU,5023
 yggdrasil/types/cast/spark_polars_cast.py,sha256=btmZNHXn2NSt3fUuB4xg7coaE0RezIBdZD92H8NK0Jw,9073
-ygg-0.1.56.dist-info/METADATA,sha256=Rr4DBB8q39XEEzBTQtYB_pja448GR74cEh3gc8-BKvY,18528
-ygg-0.1.56.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
-ygg-0.1.56.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
-ygg-0.1.56.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
-ygg-0.1.56.dist-info/RECORD,,
+ygg-0.1.57.dist-info/METADATA,sha256=0VEcri5fh3BUYJxxhQ_icTZfhkry6KgEhbKQEEDrKJ4,18528
+ygg-0.1.57.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+ygg-0.1.57.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
+ygg-0.1.57.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.57.dist-info/RECORD,,

{ygg-0.1.56.dist-info → ygg-0.1.57.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.10.1)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

yggdrasil/databricks/ai/__init__.py

@@ -1 +0,0 @@
-from .loki import *

yggdrasil/databricks/ai/loki.py

@@ -1,374 +1,53 @@
-"""
-loki.py
-
-Databricks Model Serving (OpenAI-compatible) wrapper with:
-- Loki.ask(): stateless call
-- TradingChatSession: stateful commodity trading analytics chat
-- SqlChatSession: stateful Databricks SQL generator chat
-
-Important constraint:
-- Gemini models only support ONE system prompt.
-  => We must NOT send multiple system messages.
-  => We fold summary + context blocks into a single system string.
-"""
-
-from __future__ import annotations
-
-import json
-from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Union
+from typing import Optional
+from dataclasses import field, dataclass
 
 from ..workspaces.workspace import WorkspaceService
 
 try:
-    from openai import OpenAI as _OpenAI  # noqa: F401
-except ImportError:
-    _OpenAI = None  # type: ignore
-
-
-def make_ai_client(api_key: str, base_url: str):
-    """Late import so module can load even if openai isn't installed."""
     from openai import OpenAI
-    return OpenAI(api_key=api_key, base_url=base_url)
-
-
-__all__ = ["Loki", "TradingChatSession", "SqlChatSession"]
-
-
-DEFAULT_TRADING_SYSTEM_PROMPT = """You are Loki: a conversational commodity trading analytics copilot.
 
-Scope:
-- Commodity trading analytics: curves, forwards, spreads, basis, hedging, risk, PnL explain, inventory, scheduling.
-- Databricks-first workflows: Spark/Delta/Unity Catalog, Databricks SQL, performant Python.
-
-Rules:
-- Do NOT invent prices, positions, PnL, risk, or market facts not provided.
-- State assumptions explicitly (units, time conventions, delivery months, calendars).
-- Prefer actionable output (SQL + efficient Python). Avoid slow patterns.
-- If data is missing, list exactly what you need and proceed with a reasonable template.
-
-Style:
-- Be concise, practical, performance-focused.
-"""
-
-DEFAULT_SQL_SYSTEM_PROMPT = """You are LokiSQL: a Databricks SQL generator for commodity trading analytics.
-
-Hard rules:
-- Output ONLY SQL unless the user explicitly asks for explanation.
-- Use Databricks SQL / Spark SQL dialect.
-- Prefer readable CTEs, explicit column lists, deterministic joins.
-- Do NOT invent table/column names. If missing, use placeholders like <table>, <col>.
-- If ambiguous, output best-effort SQL template with SQL comments (-- TODO ...) and placeholders.
-- Performance: push filters early, avoid exploding joins, avoid SELECT *.
+    def make_openai_client(
+        api_key: str,
+        base_url: str
+    ):
+        return OpenAI(
+            api_key=api_key,
+            base_url=base_url
+        )
+except ImportError:
+    class OpenAI:
+        pass
+
+    def make_openai_client(
+        api_key: str,
+        base_url: str
+    ):
+        from openai import OpenAI
+
+        return OpenAI(
+            api_key=api_key,
+            base_url=base_url
+        )
 
-Default assumptions:
-- Dates UTC unless specified.
-"""
+__all__ = [
+    "Loki"
+]
 
 
 @dataclass
 class Loki(WorkspaceService):
-    """
-    Loki wraps an OpenAI-compatible client pointing at Databricks Model Serving endpoints.
-    """
-
     model: str = "databricks-gemini-2-5-flash"
-    _ai_client: Optional[Any] = field(repr=False, hash=False, default=None)
+
+    _openai_client: Optional[OpenAI] = field(repr=False, hash=False, default=None)
 
     @property
-    def ai_client(self):
-        if self._ai_client is None:
-            self._ai_client = self.make_aiclient()
-        return self._ai_client
+    def openai_client(self):
+        if self._openai_client is None:
+            self._openai_client = self.make_openai_client()
+        return self._openai_client
 
-    def make_aiclient(self):
-        host = self.workspace.host.rstrip("/")
-        return make_ai_client(
+    def make_openai_client(self):
+        return make_openai_client(
             api_key=self.workspace.current_token(),
-            base_url=f"{host}/serving-endpoints",
-        )
-
-    def ask(
-        self,
-        command: str,
-        *,
-        system: Optional[str] = None,
-        max_tokens: int = 5000,
-        temperature: Optional[float] = None,
-        extra_messages: Optional[List[Dict[str, str]]] = None,
-        **kwargs: Any,
-    ) -> str:
-        """
-        Stateless single call to the model.
-
-        NOTE (Gemini constraint):
-        - Provide at most ONE system prompt (i.e., a single system message).
-        - Do not pass additional messages with role="system".
-        """
-        messages: List[Dict[str, str]] = []
-        if system:
-            messages.append({"role": "system", "content": system})
-
-        if extra_messages:
-            # IMPORTANT: caller must not include any "system" roles here for Gemini models
-            messages.extend(extra_messages)
-
-        messages.append({"role": "user", "content": command})
-
-        params: Dict[str, Any] = dict(
-            model=self.model,
-            messages=messages,
-            max_tokens=max_tokens,
-            **kwargs,
-        )
-        if temperature is not None:
-            params["temperature"] = temperature
-
-        resp = self.ai_client.chat.completions.create(**params)
-        return resp.choices[0].message.content or ""
-
-    def new_trading_chat(
-        self,
-        *,
-        system_prompt: str = DEFAULT_TRADING_SYSTEM_PROMPT,
-        max_context_turns: int = 20,
-        max_context_chars: int = 120_000,
-    ) -> "TradingChatSession":
-        return TradingChatSession(
-            loki=self,
-            system_prompt=system_prompt,
-            max_context_turns=max_context_turns,
-            max_context_chars=max_context_chars,
-        )
-
-    def new_sql_chat(
-        self,
-        *,
-        system_prompt: str = DEFAULT_SQL_SYSTEM_PROMPT,
-        max_context_turns: int = 20,
-        max_context_chars: int = 120_000,
-    ) -> "SqlChatSession":
-        return SqlChatSession(
-            loki=self,
-            system_prompt=system_prompt,
-            max_context_turns=max_context_turns,
-            max_context_chars=max_context_chars,
-        )
-
-
-@dataclass
-class _BaseChatSession:
-    """
-    Stateful session that maintains history + injected context blocks.
-
-    Gemini constraint:
-    - We must fold ALL system content into one system string.
-    - Therefore summary/context_blocks are concatenated into the system prompt.
-    """
-    loki: Loki
-    system_prompt: str
-
-    history: List[Dict[str, str]] = field(default_factory=list)
-    summary: Optional[str] = None
-    context_blocks: List[str] = field(default_factory=list)
-
-    max_context_turns: int = 20
-    max_context_chars: int = 120_000
-
-    def reset(self) -> None:
-        self.history.clear()
-        self.summary = None
-        self.context_blocks.clear()
-
-    def add_context(self, title: str, payload: Union[str, Dict[str, Any], List[Any]]) -> None:
-        if isinstance(payload, str):
-            payload_str = payload
-        else:
-            payload_str = json.dumps(payload, ensure_ascii=False, indent=2)
-
-        self.context_blocks.append(f"[Context: {title}]\n{payload_str}".strip())
-        self._trim()
-
-    def _estimate_chars(self, msgs: List[Dict[str, str]]) -> int:
-        return sum(len(m.get("content", "")) for m in msgs)
-
-    def _build_system(self, extra_system: Optional[str] = None) -> str:
-        parts: List[str] = [self.system_prompt.strip()]
-        if extra_system:
-            parts.append(extra_system.strip())
-        if self.summary:
-            parts.append(f"[ConversationSummary]\n{self.summary}".strip())
-        if self.context_blocks:
-            parts.append("\n\n".join(self.context_blocks).strip())
-        return "\n\n".join(p for p in parts if p)
-
-    def _trim(self) -> None:
-        # Turn trim (keep last N turns => N*2 messages)
-        if self.max_context_turns > 0:
-            max_msgs = self.max_context_turns * 2
-            if len(self.history) > max_msgs:
-                self.history = self.history[-max_msgs:]
-
-        # Char trim: shrink history first, then context blocks if needed
-        def total_chars() -> int:
-            sys_len = len(self._build_system())
-            return sys_len + self._estimate_chars(self.history)
-
-        while total_chars() > self.max_context_chars and self.history:
-            self.history = self.history[1:]
-
-        while total_chars() > self.max_context_chars and self.context_blocks:
-            self.context_blocks = self.context_blocks[1:]
-
-
-@dataclass
-class TradingChatSession(_BaseChatSession):
-    """
-    Commodity trading analytics chat session.
-    Optionally returns structured JSON for downstream automation.
-    """
-
-    def chat(
-        self,
-        user_text: str,
-        *,
-        structured: bool = True,
-        max_tokens: int = 12000,
-        temperature: Optional[float] = None,
-        **kwargs: Any,
-    ) -> Union[str, Dict[str, Any]]:
-        self._trim()
-
-        extra_system = None
-        if structured:
-            extra_system = (
-                "Respond ONLY as valid JSON with keys: "
-                "final_answer (string), assumptions (array of strings), data_needed (array of strings), "
-                "sql (string or null), python (string or null). "
-                "No markdown. No extra keys."
-            )
-
-        system = self._build_system(extra_system=extra_system)
-
-        assistant_text = self.loki.ask(
-            user_text,
-            system=system,
-            extra_messages=self.history,  # NOTE: history must contain no system roles
-            max_tokens=max_tokens,
-            temperature=temperature,
-            **kwargs,
+            base_url=self.workspace.host + "/serving-endpoints"
         )
-
-        self.history.append({"role": "user", "content": user_text})
-        self.history.append({"role": "assistant", "content": assistant_text})
-        self._trim()
-
-        if structured:
-            parsed = _try_parse_json_object(assistant_text)
-            if parsed is not None:
-                return parsed
-
-        return assistant_text
-
-
-@dataclass
-class SqlChatSession(_BaseChatSession):
-    """
-    SQL-only conversational session that generates Databricks SQL.
-
-    Uses a single system message with strict instructions to output SQL only.
-    """
-
-    def generate_sql(
-        self,
-        request: str,
-        *,
-        max_tokens: int = 12000,
-        temperature: Optional[float] = None,
-        sql_only: bool = True,
-        **kwargs: Any,
-    ) -> str:
-        self._trim()
-
-        extra_system = None
-        if sql_only:
-            extra_system = (
-                "Reminder: Output ONLY SQL. "
-                "If ambiguity exists, use SQL comments (-- TODO ...) and placeholders, but still output SQL only."
-            )
-
-        system = self._build_system(extra_system=extra_system)
-
-        sql = self.loki.ask(
-            request,
-            system=system,
-            extra_messages=self.history,  # history must contain no system roles
-            max_tokens=max_tokens,
-            temperature=temperature,
-            **kwargs,
-        ).strip()
-
-        self.history.append({"role": "user", "content": request})
-        self.history.append({"role": "assistant", "content": sql})
-        self._trim()
-
-        return _strip_sql_fences(sql)
-
-
-def _strip_sql_fences(text: str) -> str:
-    t = text.strip()
-    if t.startswith("```"):
-        lines = t.splitlines()
-        lines = lines[1:]  # drop ``` or ```sql
-        if lines and lines[-1].strip().startswith("```"):
-            lines = lines[:-1]
-        return "\n".join(lines).strip()
-    return t
-
-
-def _strip_markdown_fences(text: str) -> str:
-    """
-    Remove ```lang ... ``` fences if present.
-    Keeps inner content unchanged.
-    """
-    t = text.strip()
-    if not t.startswith("```"):
-        return t
-
-    lines = t.splitlines()
-    if not lines:
-        return t
-
-    # Drop first line: ``` or ```json
-    lines = lines[1:]
-
-    # Drop last line if it's ```
-    if lines and lines[-1].strip().startswith("```"):
-        lines = lines[:-1]
-
-    return "\n".join(lines).strip()
-
-
-def _try_parse_json_object(text: str) -> Optional[Dict[str, Any]]:
-    t = _strip_markdown_fences(text).strip()
-
-    # Best effort extraction if there's extra junk around JSON
-    if not t.startswith("{"):
-        start = t.find("{")
-        end = t.rfind("}")
-        if start != -1 and end != -1 and end > start:
-            t = t[start : end + 1]
-
-    try:
-        obj = json.loads(t)
-    except Exception:
-        return None
-
-    if not isinstance(obj, dict):
-        return None
-
-    required = {"final_answer", "assumptions", "data_needed", "sql", "python"}
-    if not required.issubset(set(obj.keys())):
-        return None
-
-    return obj

yggdrasil/databricks/compute/execution_context.py

@@ -16,7 +16,6 @@ from threading import Thread
 from types import ModuleType
 from typing import TYPE_CHECKING, Optional, Any, Callable, List, Dict, Union, Iterable, Tuple
 
-from .exceptions import CommandAborted
 from ...libs.databrickslib import databricks_sdk
 from ...pyutils.exceptions import raise_parsed_traceback
 from ...pyutils.expiring_dict import ExpiringDict
@@ -111,12 +110,16 @@ class ExecutionContext:
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Exit the context manager and close the remote context if created."""
         if not self._was_connected:
-            self.close(wait=False)
+            self.close()
         self.cluster.__exit__(exc_type, exc_val=exc_val, exc_tb=exc_tb)
 
     def __del__(self):
         """Best-effort cleanup for the remote execution context."""
-        self.close(wait=False)
+        if self.context_id:
+            try:
+                Thread(target=self.close).start()
+            except BaseException:
+                pass
 
     @property
     def remote_metadata(self) -> RemoteMetadata:
@@ -177,7 +180,7 @@ print(json.dumps(meta))"""
         """
         return self.cluster.workspace.sdk()
 
-    def create(
+    def create_command(
         self,
         language: "Language",
     ) -> any:
@@ -194,17 +197,15 @@ print(json.dumps(meta))"""
             self.cluster
         )
 
-        client = self._workspace_client().command_execution
-
         try:
-            created = client.create_and_wait(
+            created = self._workspace_client().command_execution.create_and_wait(
                 cluster_id=self.cluster.cluster_id,
                 language=language,
             )
         except:
             self.cluster.ensure_running()
 
-            created = client.create_and_wait(
+            created = self._workspace_client().command_execution.create_and_wait(
                 cluster_id=self.cluster.cluster_id,
                 language=language,
             )
@@ -216,38 +217,42 @@ print(json.dumps(meta))"""
 
         created = getattr(created, "response", created)
 
-        self.context_id = created.id
-
-        return self
+        return created
 
     def connect(
         self,
-        language: Optional["Language"] = None,
-        reset: bool = False
+        language: Optional["Language"] = None
     ) -> "ExecutionContext":
         """Create a remote command execution context if not already open.
 
         Args:
             language: Optional language override for the context.
-            reset: Reset existing if connected
 
         Returns:
             The connected ExecutionContext instance.
         """
         if self.context_id is not None:
-            if not reset:
-                return self
+            return self
 
-            self.close(wait=False)
+        self.language = language or self.language
 
-        language = language or self.language
+        if self.language is None:
+            self.language = Language.PYTHON
 
-        if language is None:
-            language = Language.PYTHON
+        ctx = self.create_command(language=self.language)
 
-        return self.create(language=language)
+        context_id = ctx.id
+        if not context_id:
+            raise RuntimeError("Failed to create command execution context")
 
-    def close(self, wait: bool = True) -> None:
+        self.context_id = context_id
+        LOGGER.info(
+            "Opened execution context for %s",
+            self
+        )
+        return self
+
+    def close(self) -> None:
         """Destroy the remote command execution context if it exists.
 
         Returns:
@@ -256,23 +261,12 @@ print(json.dumps(meta))"""
         if not self.context_id:
             return
 
-        client = self._workspace_client()
-
         try:
-            if wait:
-                client.command_execution.destroy(
-                    cluster_id=self.cluster.cluster_id,
-                    context_id=self.context_id,
-                )
-            else:
-                Thread(
-                    target=client.command_execution.destroy,
-                    kwargs={
-                        "cluster_id": self.cluster.cluster_id,
-                        "context_id": self.context_id,
-                    }
-                ).start()
-        except BaseException:
+            self._workspace_client().command_execution.destroy(
+                cluster_id=self.cluster.cluster_id,
+                context_id=self.context_id,
+            )
+        except Exception:
             # non-fatal: context cleanup best-effort
             pass
         finally:
@@ -471,18 +465,7 @@ print(json.dumps(meta))"""
         )
 
         try:
-            return self._decode_result(
-                result,
-                result_tag=result_tag,
-                print_stdout=print_stdout
-            )
-        except CommandAborted:
-            return self.connect(language=self.language, reset=True).execute_command(
-                command=command,
-                timeout=timeout,
-                result_tag=result_tag,
-                print_stdout=print_stdout
-            )
+            return self._decode_result(result, result_tag=result_tag, print_stdout=print_stdout)
         except ModuleNotFoundError as remote_module_error:
             _MOD_NOT_FOUND_RE = re.compile(r"No module named ['\"]([^'\"]+)['\"]")
             module_name = _MOD_NOT_FOUND_RE.search(str(remote_module_error))
@@ -677,9 +660,6 @@ with zipfile.ZipFile(buf, "r") as zf:
         if res.result_type == ResultType.ERROR:
             message = res.cause or "Command execution failed"
 
-            if "client terminated the session" in message:
-                raise CommandAborted(message)
-
             if self.language == Language.PYTHON:
                 raise_parsed_traceback(message)
 
@@ -688,7 +668,6 @@ with zipfile.ZipFile(buf, "r") as zf:
                 or getattr(res, "stack_trace", None)
                 or getattr(res, "traceback", None)
             )
-
             if remote_tb:
                 message = f"{message}\n{remote_tb}"
 

yggdrasil/databricks/workspaces/io.py

@@ -20,6 +20,7 @@ from .path_kind import DatabricksPathKind
 from ...libs.databrickslib import databricks
 from ...libs.pandaslib import PandasDataFrame
 from ...libs.polarslib import polars, PolarsDataFrame
+from ...pyutils import retry
 from ...types.cast.registry import convert
 from ...types.file_format import ExcelFileFormat
 
@@ -30,6 +31,7 @@ if databricks is not None:
         ResourceDoesNotExist,
         BadRequest,
     )
+    from databricks.sdk.errors import InternalError
 
 if TYPE_CHECKING:
     from .path import DatabricksPath
@@ -1036,6 +1038,7 @@ class DatabricksVolumeIO(DatabricksIO):
         end = start + length
         return data[start:end]
 
+    @retry(exceptions=(InternalError,))
     def write_all_bytes(self, data: bytes):
         """Write bytes to a volume file.
 

yggdrasil/databricks/workspaces/path.py

@@ -30,7 +30,7 @@ from ...types.cast.registry import convert, register_converter
 from ...types.file_format import ExcelFileFormat
 
 if databricks is not None:
-    from databricks.sdk.service.catalog import VolumeType, VolumeInfo
+    from databricks.sdk.service.catalog import VolumeType, PathOperation, VolumeInfo
     from databricks.sdk.service.workspace import ObjectType
     from databricks.sdk.errors.platform import (
         NotFound,
@@ -1236,8 +1236,6 @@ class DatabricksPath:
         self,
         operation: Optional["PathOperation"] = None
     ):
-        from databricks.sdk.service.catalog import PathOperation
-
         if self.kind != DatabricksPathKind.VOLUME:
             raise ValueError(f"Cannot generate temporary credentials for {repr(self)}")
 

yggdrasil/databricks/workspaces/workspace.py

@@ -652,9 +652,9 @@ class Workspace:
         """Return a Cluster helper bound to this workspace.
 
         Args:
-            workspace: Optional workspace override.
             cluster_id: Optional cluster id.
             cluster_name: Optional cluster name.
+            **kwargs: Additional Cluster parameters.
 
         Returns:
             A Cluster instance.
@@ -662,7 +662,7 @@ class Workspace:
         from ..compute.cluster import Cluster
 
         return Cluster(
-            workspace=self if workspace is None else workspace,
+            workspace=self,
             cluster_id=cluster_id,
             cluster_name=cluster_name,
         )
@@ -671,19 +671,10 @@ class Workspace:
         self,
         workspace: Optional["Workspace"] = None,
     ):
-        """
-        Return a Cluster helper bound to this workspace.
-
-        Args:
-            workspace: Optional workspace override.
-
-        Returns:
-            A Loki AI instance.
-        """
         from ..ai.loki import Loki
 
         return Loki(
-            workspace=self if workspace is None else workspace,
+            workspace=self,
         )
 
 # ---------------------------------------------------------------------------

yggdrasil/pyutils/__init__.py

@@ -4,3 +4,5 @@ from .retry import retry
 from .parallel import parallelize
 from .python_env import PythonEnv
 from .callable_serde import CallableSerde
+
+__all__ = ["retry", "parallelize", "PythonEnv", "CallableSerde"]

yggdrasil/pyutils/python_env.py

@@ -20,7 +20,7 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Iterable, Iterator, Mapping, MutableMapping, Optional, Union, List, Tuple
 
-from .modules import PipIndexSettings
+from yggdrasil.pyutils.modules import PipIndexSettings
 
 log = logging.getLogger(__name__)
 

yggdrasil/version.py

@@ -1 +1 @@
-__version__ = "0.1.56"
+__version__ = "0.1.57"

yggdrasil/databricks/compute/exceptions.py

@@ -1,14 +0,0 @@
-from ...exceptions import YGGException
-
-__all__ = [
-    "ComputeException",
-    "CommandAborted"
-]
-
-
-class ComputeException(YGGException):
-    pass
-
-
-class CommandAborted(YGGException):
-    pass

yggdrasil/exceptions.py

@@ -1,7 +0,0 @@
-__all__ = [
-    "YGGException"
-]
-
-
-class YGGException(Exception):
-    pass