ygg 0.1.16__tar.gz → 0.1.18__tar.gz

{ygg-0.1.16 → ygg-0.1.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.16
+Version: 0.1.18
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 Project-URL: Homepage, https://github.com/Platob/Yggdrasil

{ygg-0.1.16 → ygg-0.1.18}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ygg"
-version = "0.1.16"
+version = "0.1.18"
 description = "Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks"
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { file = "LICENSE" }

{ygg-0.1.16 → ygg-0.1.18}/src/ygg.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.16
+Version: 0.1.18
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 Project-URL: Homepage, https://github.com/Platob/Yggdrasil

{ygg-0.1.16 → ygg-0.1.18}/src/ygg.egg-info/SOURCES.txt

@@ -40,7 +40,7 @@ src/yggdrasil/requests/__init__.py
 src/yggdrasil/requests/msal.py
 src/yggdrasil/requests/session.py
 src/yggdrasil/ser/__init__.py
-src/yggdrasil/ser/method.py
+src/yggdrasil/ser/callable_serde.py
 src/yggdrasil/types/__init__.py
 src/yggdrasil/types/libs.py
 src/yggdrasil/types/python_arrow.py

{ygg-0.1.16 → ygg-0.1.18}/src/yggdrasil/databricks/compute/cluster.py

@@ -26,7 +26,7 @@ from ..workspaces.workspace import WorkspaceService, Workspace
 from ... import retry
 from ...libs.databrickslib import databricks_sdk
 from ...pyutils.modules import PipIndexSettings
-from ...ser import SerializedFunction
+from ...ser import CallableSerdeMixin
 
 if databricks_sdk is None:  # pragma: no cover - import guard
     ResourceDoesNotExist = Exception  # type: ignore
@@ -154,10 +154,17 @@ class Cluster(WorkspaceService):
             ] if _ not in libraries
         ])
 
+        python_version = sys.version_info
+
+        if python_version[0] < 3:
+            python_version = None
+        elif python_version[1] < 11:
+            python_version = None
+
         inst = self.create_or_update(
             cluster_id=cluster_id,
             cluster_name=cluster_name or self.cluster_name or self.workspace.current_user.user_name,
-            python_version=sys.version_info,
+            python_version=python_version,
             single_user_name=single_user_name or self.workspace.current_user.user_name,
             runtime_engine=runtime_engine or RuntimeEngine.PHOTON,
             libraries=libraries,
@@ -172,8 +179,11 @@ class Cluster(WorkspaceService):
             self.details = self.clusters_client().get(cluster_id=self.cluster_id)
         return self._details
 
-    def fresh_details(self, max_delay: float):
-        if self.cluster_id and time.time() - self._details_refresh_time > max_delay:
+    def fresh_details(self, max_delay: float | None = None):
+        max_delay = max_delay or 0
+        delay = time.time() - self._details_refresh_time
+
+        if self.cluster_id and delay > max_delay:
             self.details = self.clusters_client().get(cluster_id=self.cluster_id)
         return self._details
 
@@ -187,14 +197,59 @@ class Cluster(WorkspaceService):
 
     @property
     def state(self):
-        if self.cluster_id:
-            return self.fresh_details(max_delay=10).state
+        details = self.fresh_details(max_delay=10)
+
+        if details is not None:
+            return details.state
+        return State.UNKNOWN
+
+    def get_state(self, max_delay: float = None):
+        details = self.fresh_details(max_delay=max_delay)
+
+        if details is not None:
+            return details.state
         return State.UNKNOWN
 
     @property
     def is_running(self):
         return self.state == State.RUNNING
 
+    @property
+    def is_pending(self):
+        return self.state in  (State.PENDING, State.RESIZING, State.RESTARTING, State.TERMINATING)
+
+    @property
+    def is_error(self):
+        return self.state == State.ERROR
+
+    def raise_for_status(self):
+        if self.is_error:
+            raise DatabricksError("Error in %s" % self)
+
+        return self
+
+    def wait_for_status(
+        self,
+        tick: float = 0.5,
+        timeout: float = 600,
+        backoff: int = 2,
+        max_sleep_time: float = 15
+    ):
+        start = time.time()
+        sleep_time = tick
+
+        while self.is_pending:
+            time.sleep(sleep_time)
+
+            if time.time() - start > timeout:
+                raise TimeoutError("Waiting state for %s timed out")
+
+            sleep_time = min(max_sleep_time, sleep_time * backoff)
+
+        self.raise_for_status()
+
+        return self
+
     @property
     def spark_version(self) -> str:
         d = self.details
@@ -276,7 +331,8 @@ class Cluster(WorkspaceService):
 
             versions = [v for v in versions if py_for_key(v.key) == py_filter]
 
-            if not versions and py_filter > 12:
+            # Handle superior pyton versions
+            if not versions and py_filter[1] > 12:
                 return self.spark_versions(photon=photon)
 
         return versions
@@ -402,7 +458,7 @@ class Cluster(WorkspaceService):
         libraries: Optional[List[Union[str, "Library"]]] = None,
         **cluster_spec: Any
     ) -> "Cluster":
-        self.install_libraries(libraries=libraries, timeout=None, raise_error=False)
+        self.install_libraries(libraries=libraries, wait_timeout=None, raise_error=False)
 
         existing_details = {
             k: v
@@ -482,24 +538,29 @@ class Cluster(WorkspaceService):
     ) -> "Cluster":
         return self.start()
 
-    @retry(tries=4)
     def start(
         self,
     ) -> "Cluster":
+        self.wait_for_status()
+
         if not self.is_running:
             logger.info("Starting %s", self)
             self.details = self.clusters_client().start_and_wait(cluster_id=self.cluster_id)
             return self.wait_installed_libraries()
+
         return self
 
     @retry(tries=4)
     def restart(
         self,
     ):
+        self.wait_for_status()
+
         if self.is_running:
             logger.info("Restarting %s", self)
             self.details = self.clusters_client().restart_and_wait(cluster_id=self.cluster_id)
             return self.wait_installed_libraries()
+
         return self.start()
 
     def delete(
@@ -530,15 +591,14 @@ class Cluster(WorkspaceService):
         timeout: Optional[dt.timedelta] = None,
         result_tag: Optional[str] = None,
     ):
-        with self.execution_context(language=language) as ctx:
-            return ctx.execute(
-                obj=obj,
-                args=args,
-                kwargs=kwargs,
-                env_keys=env_keys,
-                timeout=timeout,
-                result_tag=result_tag
-            )
+        return self.execution_context(language=language).execute(
+            obj=obj,
+            args=args,
+            kwargs=kwargs,
+            env_keys=env_keys,
+            timeout=timeout,
+            result_tag=result_tag
+        )
 
     # ------------------------------------------------------------------
     # decorator that routes function calls via `execute`
@@ -551,6 +611,7 @@ class Cluster(WorkspaceService):
         env_keys: Optional[List[str]] = None,
         timeout: Optional[dt.timedelta] = None,
         result_tag: Optional[str] = None,
+        **options
     ):
         """
         Decorator to run a function via Workspace.execute instead of locally.
@@ -570,7 +631,7 @@ class Cluster(WorkspaceService):
         """
         def decorator(func: Callable):
             context = self.execution_context(language=language or Language.PYTHON)
-            serialized = func if isinstance(func, SerializedFunction) else SerializedFunction.from_callable(func)
+            serialized = CallableSerdeMixin.from_callable(func)
 
             @functools.wraps(func)
             def wrapper(*args, **kwargs):
@@ -584,6 +645,7 @@ class Cluster(WorkspaceService):
                     env_keys=env_keys,
                     timeout=timeout,
                     result_tag=result_tag,
+                    **options
                 )
 
             return wrapper
@@ -597,39 +659,96 @@ class Cluster(WorkspaceService):
     def install_libraries(
         self,
         libraries: Optional[List[Union[str, "Library"]]] = None,
-        timeout: Optional[dt.timedelta] = dt.timedelta(minutes=5),
+        wait_timeout: Optional[dt.timedelta] = dt.timedelta(minutes=20),
         pip_settings: Optional[PipIndexSettings] = None,
-        raise_error: bool = True
+        raise_error: bool = True,
+        restart: bool = True,
     ) -> "Cluster":
         if not libraries:
             return self
 
         wsdk = self.workspace.sdk()
 
-        wsdk.libraries.install(
-            cluster_id=self.cluster_id,
-            libraries=[
-                self._check_library(_, pip_settings=pip_settings)
-                for _ in libraries if _
+        libraries = [
+            self._check_library(_, pip_settings=pip_settings)
+            for _ in libraries if _
+        ]
+
+        if libraries:
+            wsdk.libraries.install(
+                cluster_id=self.cluster_id,
+                libraries=[
+                    self._check_library(_, pip_settings=pip_settings)
+                    for _ in libraries if _
+                ]
+            )
+
+            if wait_timeout is not None:
+                self.wait_installed_libraries(
+                    timeout=wait_timeout, pip_settings=pip_settings, raise_error=raise_error
+                )
+
+        return self
+
+    def installed_library_statuses(self):
+        return self.workspace.sdk().libraries.cluster_status(cluster_id=self.cluster_id)
+
+    def uninstall_libraries(
+        self,
+        pypi_packages: Optional[list[str]] = None,
+        libraries: Optional[list["Library"]] = None,
+        restart: bool = True
+    ):
+        if libraries is None:
+            to_remove = [
+                lib.library
+                for lib in self.installed_library_statuses()
+                if self._filter_lib(
+                    lib,
+                    pypi_packages=pypi_packages,
+                    default_filter=False
+                )
             ]
-        )
+        else:
+            to_remove = libraries
+
+        if to_remove:
+            self.workspace.sdk().libraries.uninstall(
+                cluster_id=self.cluster_id,
+                libraries=to_remove
+            )
 
-        if timeout is not None:
-            self.wait_installed_libraries(timeout=timeout, pip_settings=pip_settings, raise_error=raise_error)
+            if restart:
+                self.restart()
 
         return self
 
+    @staticmethod
+    def _filter_lib(
+        lib: Optional["Library"],
+        pypi_packages: Optional[list[str]] = None,
+        default_filter: bool = False
+    ):
+        if lib is None:
+            return False
+
+        if lib.pypi:
+            if lib.pypi.package and pypi_packages:
+                return lib.pypi.package in pypi_packages
+
+        return default_filter
+
     def wait_installed_libraries(
         self,
-        timeout: dt.timedelta = dt.timedelta(minutes=5),
+        timeout: dt.timedelta = dt.timedelta(minutes=20),
         pip_settings: Optional[PipIndexSettings] = None,
-        raise_error: bool = True
+        raise_error: bool = True,
     ):
         if not self.is_running:
             return self
 
-        wsdk = self.workspace.sdk()
-        statuses = list(wsdk.libraries.cluster_status(cluster_id=self.cluster_id))
+        statuses = list(self.installed_library_statuses())
+
         max_time = time.time() + timeout.total_seconds()
 
         while True:
@@ -641,6 +760,7 @@ class Cluster(WorkspaceService):
             if failed:
                 if raise_error:
                     raise DatabricksError("Libraries %s in %s failed to install" % (failed, self))
+
                 logger.warning(
                     "Libraries %s in %s failed to install",
                     failed, self
@@ -662,7 +782,7 @@ class Cluster(WorkspaceService):
                 )
 
             time.sleep(10)
-            statuses = list(wsdk.libraries.cluster_status(cluster_id=self.cluster_id))
+            statuses = list(self.installed_library_statuses())
 
         return self
 
@@ -680,6 +800,8 @@ class Cluster(WorkspaceService):
         if isinstance(value, Library):
             return value
 
+        pip_settings = PipIndexSettings.default_settings() if pip_settings is None else pip_settings
+
         if isinstance(value, str):
             if os.path.exists(value):
                 target_path = self.workspace.shared_cache_path(
@@ -698,11 +820,11 @@ class Cluster(WorkspaceService):
             elif value.endswith(".whl"):
                 return Library(whl=value)
 
-            # Fallback: treat as PyPI / private index package
-            if pip_settings:
-                repo = pip_settings.extra_index_urls[0] if pip_settings.extra_index_urls else None
-            else:
-                repo = None
+            repo = None
+
+            if pip_settings.extra_index_url:
+                if value.startswith("datamanagement") or value.startswith("TSSecrets") or value.startswith("tgp_"):
+                    repo = pip_settings.extra_index_url
 
             return Library(
                 pypi=PythonPyPiLibrary(