ygg 0.1.20__py3-none-any.whl → 0.1.23__py3-none-any.whl

yggdrasil/databricks/compute/cluster.py

@@ -15,7 +15,6 @@ import functools
 import inspect
 import logging
 import os
-import sys
 import time
 from dataclasses import dataclass
 from types import ModuleType
@@ -23,10 +22,10 @@ from typing import Any, Iterator, Optional, Union, List, Callable, Dict, ClassVa
 
 from .execution_context import ExecutionContext
 from ..workspaces.workspace import WorkspaceService, Workspace
-from ... import retry
+from ... import retry, CallableSerde
 from ...libs.databrickslib import databricks_sdk
 from ...pyutils.modules import PipIndexSettings
-from ...ser import CallableSerdeMixin
+from ...pyutils.python_env import PythonEnv
 
 if databricks_sdk is None:  # pragma: no cover - import guard
     ResourceDoesNotExist = Exception  # type: ignore
@@ -127,7 +126,7 @@ class Cluster(WorkspaceService):
         # 🔥 first time for this host → create
         inst = cls._env_clusters[host] = (
             cls(workspace=workspace, cluster_id=cluster_id, cluster_name=cluster_name)
-            .replicate_current_environment(
+            .push_python_environment(
                 single_user_name=single_user_name,
                 runtime_engine=runtime_engine,
                 libraries=libraries,
@@ -137,8 +136,9 @@ class Cluster(WorkspaceService):
 
         return inst
     
-    def replicate_current_environment(
+    def push_python_environment(
         self,
+        source: Optional[PythonEnv] = None,
         cluster_id: Optional[str] = None,
         cluster_name: Optional[str] = None,
         single_user_name: Optional[str] = None,
@@ -146,15 +146,19 @@ class Cluster(WorkspaceService):
         libraries: Optional[list[str]] = None,
         **kwargs
     ) -> "Cluster":
+        if source is None:
+            source = PythonEnv.get_current()
+
         libraries = list(libraries) if libraries is not None else []
         libraries.extend([
             _ for _ in [
                 "ygg",
-                "dill"
+                "dill",
+                "uv",
             ] if _ not in libraries
         ])
 
-        python_version = sys.version_info
+        python_version = source.version_info
 
         if python_version[0] < 3:
             python_version = None
@@ -172,6 +176,28 @@ class Cluster(WorkspaceService):
         )
 
         return inst
+
+    def pull_python_environment(
+        self,
+        target: Optional[PythonEnv] = None,
+    ):
+        with self.context() as c:
+            m = c.remote_metadata
+            requirements = m.requirements
+            version_info = m.version_info
+
+        if target is None:
+            target = PythonEnv.create(
+                name=f"dbx-{self.name}",
+                python=".".join(str(_) for _ in version_info)
+            )
+        else:
+            target.update(
+                requirements=requirements,
+                python=".".join(str(_) for _ in version_info)
+            )
+
+        return target
     
     @property
     def details(self):
@@ -364,6 +390,8 @@ class Cluster(WorkspaceService):
         python_version: Optional[Union[str, tuple[int, ...]]] = None,
         **kwargs
     ):
+        pip_settings = PipIndexSettings.default_settings()
+
         if kwargs:
             details = ClusterDetails(**{
                 **details.as_shallow_dict(),
@@ -395,6 +423,13 @@ class Cluster(WorkspaceService):
         if details.is_single_node is not None and details.kind is None:
             details.kind = Kind.CLASSIC_PREVIEW
 
+        if pip_settings.extra_index_urls:
+            if details.spark_env_vars is None:
+                details.spark_env_vars = {}
+            str_urls = " ".join(pip_settings.extra_index_urls)
+            details.spark_env_vars["UV_EXTRA_INDEX_URL"] = details.spark_env_vars.get("UV_INDEX", str_urls)
+            details.spark_env_vars["PIP_EXTRA_INDEX_URL"] = details.spark_env_vars.get("PIP_EXTRA_INDEX_URL", str_urls)
+
         return details
 
     def create_or_update(
@@ -569,7 +604,7 @@ class Cluster(WorkspaceService):
         logger.info("Deleting %s", self)
         return self.clusters_client().delete(cluster_id=self.cluster_id)
 
-    def execution_context(
+    def context(
         self,
         language: Optional["Language"] = None,
         context_id: Optional[str] = None
@@ -591,7 +626,7 @@ class Cluster(WorkspaceService):
         timeout: Optional[dt.timedelta] = None,
         result_tag: Optional[str] = None,
     ):
-        return self.execution_context(language=language).execute(
+        return self.context(language=language).execute(
             obj=obj,
             args=args,
             kwargs=kwargs,
@@ -607,8 +642,10 @@ class Cluster(WorkspaceService):
         self,
         _func: Optional[Callable] = None,
         *,
+        before: Optional[Callable] = None,
         language: Optional["Language"] = None,
         env_keys: Optional[List[str]] = None,
+        env_variables: Optional[Dict[str, str]] = None,
         timeout: Optional[dt.timedelta] = None,
         result_tag: Optional[str] = None,
         **options
@@ -630,19 +667,23 @@ class Cluster(WorkspaceService):
             def h(z): ...
         """
         def decorator(func: Callable):
-            context = self.execution_context(language=language or Language.PYTHON)
-            serialized = CallableSerdeMixin.from_callable(func)
+            context = self.context(language=language or Language.PYTHON)
+            serialized = CallableSerde.from_callable(func)
+            do_before = CallableSerde.from_callable(before)
 
             @functools.wraps(func)
             def wrapper(*args, **kwargs):
                 if os.getenv("DATABRICKS_RUNTIME_VERSION") is not None:
                     return func(*args, **kwargs)
 
+                do_before()
+
                 return context.execute(
                     obj=serialized,
                     args=list(args),
                     kwargs=kwargs,
                     env_keys=env_keys,
+                    env_variables=env_variables,
                     timeout=timeout,
                     result_tag=result_tag,
                     **options
@@ -685,7 +726,9 @@ class Cluster(WorkspaceService):
 
             if wait_timeout is not None:
                 self.wait_installed_libraries(
-                    timeout=wait_timeout, pip_settings=pip_settings, raise_error=raise_error
+                    timeout=wait_timeout,
+                    pip_settings=pip_settings,
+                    raise_error=raise_error
                 )
 
         return self
@@ -790,7 +833,7 @@ class Cluster(WorkspaceService):
         self,
         libraries: str | ModuleType | List[str | ModuleType],
     ):
-        return self.execution_context().install_temporary_libraries(libraries=libraries)
+        return self.context().install_temporary_libraries(libraries=libraries)
 
     def _check_library(
         self,
@@ -823,7 +866,11 @@ class Cluster(WorkspaceService):
             repo = None
 
             if pip_settings.extra_index_url:
-                if value.startswith("datamanagement") or value.startswith("TSSecrets") or value.startswith("tgp_"):
+                if (
+                    value.startswith("datamanagement")
+                    or value.startswith("TSSecrets")
+                    or value.startswith("tgp_")
+                ):
                     repo = pip_settings.extra_index_url
 
             return Library(

yggdrasil/databricks/compute/execution_context.py

@@ -11,12 +11,12 @@ import sys
 import threading
 import zipfile
 from types import ModuleType
-from typing import TYPE_CHECKING, Optional, Any, Callable, List, Dict, Union, Iterable
+from typing import TYPE_CHECKING, Optional, Any, Callable, List, Dict, Union, Iterable, Tuple
 
 from ...libs.databrickslib import databricks_sdk
 from ...pyutils.exceptions import raise_parsed_traceback
 from ...pyutils.modules import resolve_local_lib_path
-from ...ser import CallableSerdeMixin
+from ...pyutils.callable_serde import CallableSerde
 
 if TYPE_CHECKING:
     from .cluster import Cluster
@@ -35,6 +35,8 @@ logger = logging.getLogger(__name__)
 class RemoteMetadata:
     site_packages_path: Optional[str] = dc.field(default=None)
     os_env: Dict[str, str] = dc.field(default_factory=dict)
+    requirements: Optional[str] = dc.field(default=None)
+    version_info: Tuple[int, int, int] = dc.field(default=(0, 0, 0))
 
     def os_env_diff(
         self,
@@ -49,6 +51,7 @@ class RemoteMetadata:
             if k not in self.os_env.keys()
         }
 
+
 @dc.dataclass
 class ExecutionContext:
     """
@@ -74,23 +77,21 @@ class ExecutionContext:
     _was_connected: Optional[bool] = None
     _remote_metadata: Optional[RemoteMetadata] = None
 
-    __lock: threading.Lock = dc.field(default_factory=threading.Lock, init=False, repr=False)
+    _lock: threading.RLock = dc.field(default_factory=threading.RLock, init=False, repr=False)
 
     # --- Pickle / cloudpickle support (don’t serialize locks or cached remote metadata) ---
     def __getstate__(self):
         state = self.__dict__.copy()
 
-        # name-mangled field for __lock in instance dict:
-        state.pop("_ExecutionContext__lock", None)
+        # name-mangled field for _lock in instance dict:
+        state.pop("_lock", None)
 
         return state
 
     def __setstate__(self, state):
-        self.__dict__.update(state)
+        state["_lock"] = state.get("_lock", threading.RLock())
 
-        # recreate lock + reset cache on unpickle
-        self.__dict__["_ExecutionContext__lock"] = threading.Lock()
-        self._remote_metadata = None
+        self.__dict__.update(state)
 
     def __enter__(self) -> "ExecutionContext":
         self.cluster.__enter__()
@@ -106,20 +107,22 @@ class ExecutionContext:
         self.close()
 
     @property
-    def remote_metadata(self) -> "RemoteMetadata":
+    def remote_metadata(self) -> RemoteMetadata:
         # fast path (no lock)
         rm = self._remote_metadata
         if rm is not None:
             return rm
 
         # slow path guarded
-        with self.__lock:
+        with self._lock:
             # double-check after acquiring lock
             if self._remote_metadata is None:
                 cmd = r"""import glob
 import json
 import os
+from yggdrasil.pyutils.python_env import PythonEnv
 
+current_env = PythonEnv.get_current()
 meta = {}
 
 for path in glob.glob('/local_**/.ephemeral_nfs/cluster_libraries/python/lib/python*/site-*', recursive=False):
@@ -130,9 +133,11 @@ for path in glob.glob('/local_**/.ephemeral_nfs/cluster_libraries/python/lib/pyt
 os_env = meta["os_env"] = {}
 for k, v in os.environ.items():
     os_env[k] = v
+    
+meta["requirements"] = current_env.export_requirements_matrix()
+meta["version_info"] = current_env.version_info
 
-print(json.dumps(meta))
-"""
+print(json.dumps(meta))"""
 
                 content = self.execute_command(
                     command=cmd,
@@ -226,7 +231,7 @@ print(json.dumps(meta))
     # ------------ public API ------------
     def execute(
         self,
-        obj: Union[str, Callable, CallableSerdeMixin],
+        obj: Union[str, Callable],
         *,
         args: List[Any] = None,
         kwargs: Dict[str, Any] = None,
@@ -260,7 +265,7 @@ print(json.dumps(meta))
 
     def execute_callable(
         self,
-        func: Callable | CallableSerdeMixin,
+        func: Callable | CallableSerde,
         args: List[Any] = None,
         kwargs: Dict[str, Any] = None,
         env_keys: Optional[Iterable[str]] = None,
@@ -269,7 +274,7 @@ print(json.dumps(meta))
         timeout: Optional[dt.timedelta] = None,
         command: Optional[str] = None,
         use_dill: Optional[bool] = None
-    ) -> str:
+    ) -> Any:
         if self.is_in_databricks_environment():
             args = args or []
             kwargs = kwargs or {}
@@ -284,7 +289,7 @@ print(json.dumps(meta))
             self,
         )
 
-        serialized = CallableSerdeMixin.from_callable(func)
+        serialized = CallableSerde.from_callable(func)
 
         self.install_temporary_libraries(libraries=serialized.package_root)
 
@@ -302,9 +307,6 @@ print(json.dumps(meta))
         command = serialized.to_command(
             args=args,
             kwargs=kwargs,
-            env_keys=env_keys or [],
-            env_variables=env_variables,
-            use_dill=use_dill,
             result_tag=result_tag,
         ) if not command else command
 

yggdrasil/databricks/compute/remote.py

@@ -9,8 +9,6 @@ from typing import (
 
 from ..workspaces.workspace import Workspace
 
-
-
 ReturnType = TypeVar("ReturnType")
 
 logger = logging.getLogger(__name__)

yggdrasil/pyutils/__init__.py

@@ -1,2 +1,4 @@
 from .retry import retry
 from .parallel import parallelize
+from .python_env import PythonEnv
+from .callable_serde import CallableSerde