ygg 0.1.16__py3-none-any.whl → 0.1.18__py3-none-any.whl

{ygg-0.1.16.dist-info → ygg-0.1.18.dist-info}/METADATA

@@ -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.dist-info → ygg-0.1.18.dist-info}/RECORD

@@ -1,9 +1,9 @@
 yggdrasil/__init__.py,sha256=6OPibApplA5TF4TeixkQO_qewpaAidYX-fSDvvKYcTI,91
 yggdrasil/databricks/__init__.py,sha256=aGVve5mpoQtxSK2nfzrexjRPoutCIyaOnKZijkG4_QE,92
 yggdrasil/databricks/compute/__init__.py,sha256=TVDwPmW2SOmHmnhzZhsvrWbrxZ_lEcgqe3l9BeB-oxM,218
-yggdrasil/databricks/compute/cluster.py,sha256=Z0igibj9osfw0sLPyO9qr1ZMXdGJAlay3xWO3fpJMms,23806
-yggdrasil/databricks/compute/execution_context.py,sha256=742neVcZWKInhDzWoqBQmFj82Hsg0AJLTFF1iB999lw,16242
-yggdrasil/databricks/compute/remote.py,sha256=NPljUmHt2ZHJmCw0EAaZAEXjTgH2QYOSO9t8-z4pRx0,1177
+yggdrasil/databricks/compute/cluster.py,sha256=xElDioObG6exkUS08K-Ccs_EFNbWD69Z15fjvnHwOx8,26958
+yggdrasil/databricks/compute/execution_context.py,sha256=pnzA_itZiYW4LpjgWnKZlovABEjHToDAp2ahVMnZmRQ,18625
+yggdrasil/databricks/compute/remote.py,sha256=DzPVPk-4bt5bOL52Onur3xLNh3UzS2K70DU5HglhGSg,1216
 yggdrasil/databricks/jobs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 yggdrasil/databricks/jobs/config.py,sha256=8Slfw4Wl7vu0kIlaUUqVqjjOgPwuULoo0rroENCbC20,11494
 yggdrasil/databricks/sql/__init__.py,sha256=JZpQ9eCphDf1l4yzIZ7a7OLigxqXkqOgb0Mio7Rj09A,181
@@ -26,14 +26,14 @@ yggdrasil/libs/extensions/polars_extensions.py,sha256=kojeJOW5VPBXMVTJT2gWUJypNu
 yggdrasil/libs/extensions/spark_extensions.py,sha256=uCcpLJr0wEE9V_-nc8DrQtAh9n2Q4mV_NicU-DNGG54,13396
 yggdrasil/pyutils/__init__.py,sha256=QF4o345TW0wRIt_SR_VuH8AnwESFtKkksnqaDurPkbY,61
 yggdrasil/pyutils/exceptions.py,sha256=Kt4xY4jPv-Ld8vBJaiQy51xbewD6Q-3bZkaa3THsFNM,2891
-yggdrasil/pyutils/modules.py,sha256=FKXUD2Swdj6fuh13yj4Tid2G1q6wjbPkzEMQGTwHXUE,10051
+yggdrasil/pyutils/modules.py,sha256=r3C4TRA8C_fUTltITV0VJY0pbIdxdV1op2z8BX0FgJo,10201
 yggdrasil/pyutils/parallel.py,sha256=L2r4_iumOzDk7omWVBPYmT5U8n7suOdv1AFpa_ghHSY,5902
 yggdrasil/pyutils/retry.py,sha256=1zjascEsffvnkVRdHhdeoenk6tBPrzrj4VhBltbhBeU,10390
 yggdrasil/requests/__init__.py,sha256=THJz1IoZYQccwmXcQR3N8D-uWxCkfMtgeXDhONdERR8,41
 yggdrasil/requests/msal.py,sha256=ucnN45iZZpbXkByw212PX4shH4g0EeyrW8JEmfimWtY,5861
 yggdrasil/requests/session.py,sha256=YomLcDf8O_mc8BUnf9fr5wrupDnxEzaGw-guhV91NsE,830
-yggdrasil/ser/__init__.py,sha256=osetghhzlABSC8dk9GDJlEwBMIhAqOTVhEDOWMHpNFg,23
-yggdrasil/ser/method.py,sha256=TljQm882NfHI9ca5Moa2W2dVTpG_JdjlnsjGSMxVtOQ,24026
+yggdrasil/ser/__init__.py,sha256=sS66Bxu8aiLb-8N2aNayquamfi7FobEH51JyV5ULDFI,31
+yggdrasil/ser/callable_serde.py,sha256=1pDgrzAceoFQ7JS7qIeSxC4hCz2oAcQOkeWyYYeT7iY,21206
 yggdrasil/types/__init__.py,sha256=p0Qu_69RkePPgQGM9nSue_bcbEIAM2u9eo3zsEplHJ8,82
 yggdrasil/types/libs.py,sha256=7-p0M4C6TnEWpUGf2nY4XshhJxtXOE_-bsYmJWU6jtk,227
 yggdrasil/types/python_arrow.py,sha256=1Ac1ZnEgmH4nLNBfrfbM6E9EfANGWtdANqlwMENQkTw,21408
@@ -48,7 +48,7 @@ yggdrasil/types/cast/registry.py,sha256=-88mq-U1pDSGbEC9PRY0zJCzloyBodXgeSRBPb6h
 yggdrasil/types/cast/spark_cast.py,sha256=IHthM78dugabGXxNNW9sSHn-olDwzXcFdIFcPo9IiXU,23021
 yggdrasil/types/cast/spark_pandas_cast.py,sha256=8PgJItF_XbyBcNuBnXkMQU3PBy3sAPEXZT9SXL2WbU4,4200
 yggdrasil/types/cast/spark_polars_cast.py,sha256=ba1UOvY1ouGCro1Np9slXmJ4TEyWnUtwVEAwxGvPLlk,8336
-ygg-0.1.16.dist-info/METADATA,sha256=2DrFcSHmSoioahWJK_YojsdaFHHuclrWmLXq3HD0Ooo,5981
-ygg-0.1.16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ygg-0.1.16.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
-ygg-0.1.16.dist-info/RECORD,,
+ygg-0.1.18.dist-info/METADATA,sha256=d_adw4j5tP0rLrz-9sgI6fL5vWp5s6KgICs7lyLctgI,5981
+ygg-0.1.18.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ygg-0.1.18.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.18.dist-info/RECORD,,

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(