ob-metaflow-extensions 1.4.5__tar.gz → 1.4.6__tar.gz

{ob_metaflow_extensions-1.4.5 → ob_metaflow_extensions-1.4.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob_metaflow_extensions
-Version: 1.4.5
+Version: 1.4.6
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob_metaflow_extensions-1.4.5 → ob_metaflow_extensions-1.4.6}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -361,4 +361,4 @@ SECRETS_PROVIDERS_DESC = [
     ("outerbounds", ".secrets.secrets.OuterboundsSecretsProvider"),
 ]
 # Adding an override here so the library can be imported at the metaflow.plugins level
-__mf_promote_submodules__ = ["snowflake", "ollama", "torchtune"]
+__mf_promote_submodules__ = ["snowflake", "ollama", "torchtune", "optuna"]

{ob_metaflow_extensions-1.4.5 → ob_metaflow_extensions-1.4.6}/metaflow_extensions/outerbounds/plugins/apps/core/capsule.py

@@ -446,6 +446,34 @@ class CapsuleApi:
                 message="Capsule JSON decode failed",
             )
 
+    def get_by_name(self, name: str, most_recent_only: bool = True):
+        _url = os.path.join(self._base_url, f"?displayName={name}")
+        response = self._wrapped_api_caller(
+            requests.get,
+            _url,
+            retryable_status_codes=[409],  # todo : verify me
+            conn_error_retries=3,
+        )
+        try:
+            if most_recent_only:
+                result = response.json()
+                candidates = result["capsules"]
+                if not candidates:
+                    return None
+                return sorted(
+                    candidates, key=lambda x: x["metadata"]["createdAt"], reverse=True
+                )[0]
+            else:
+                return response.json()
+        except json.JSONDecodeError as e:
+            raise CapsuleApiException(
+                _url,
+                "get",
+                response.status_code,
+                response.text,
+                message="Capsule JSON decode failed",
+            )
+
     def list(self):
         response = self._wrapped_api_caller(
             requests.get,
@@ -793,9 +821,11 @@ class CapsuleDeployer:
                     "💊 %s deployment status: %s "
                     % (
                         self.capsule_type.title(),
-                        "in progress"
-                        if state_machine.update_in_progress
-                        else "completed",
+                        (
+                            "in progress"
+                            if state_machine.update_in_progress
+                            else "completed"
+                        ),
                     )
                 )
                 _further_readiness_check_failed = False

{ob_metaflow_extensions-1.4.5 → ob_metaflow_extensions-1.4.6}/metaflow_extensions/outerbounds/plugins/apps/core/perimeters.py

@@ -40,7 +40,7 @@ class PerimeterExtractor:
         return perimeter, api_server  # type: ignore
 
     @classmethod
-    def during_metaflow_execution(cls) -> str:
+    def during_metaflow_execution(cls) -> Union[Tuple[str, str], Tuple[None, None]]:
         from metaflow.metaflow_config_funcs import init_config
 
         clean_url = (

ob_metaflow_extensions-1.4.6/metaflow_extensions/outerbounds/plugins/optuna/__init__.py

@@ -0,0 +1,48 @@
+import os
+import json
+
+__mf_promote_submodules__ = ["plugins.optuna"]
+
+
+def auth():
+    from metaflow.metaflow_config_funcs import init_config
+
+    conf = init_config()
+    if conf:
+        headers = {"x-api-key": conf["METAFLOW_SERVICE_AUTH_KEY"]}
+    else:
+        headers = json.loads(os.environ["METAFLOW_SERVICE_HEADERS"])
+    return headers
+
+
+def get_deployment_db_access_endpoint(name: str):
+    from ..apps.core.perimeters import PerimeterExtractor
+    from ..apps.core.capsule import CapsuleApi
+
+    perimeter, cap_url = PerimeterExtractor.during_metaflow_execution()
+    deployment = CapsuleApi(cap_url, perimeter).get_by_name(name)
+    if not deployment:
+        raise Exception(f"No app deployment found with name `{name}`")
+
+    if (
+        "status" in deployment
+        and "accessInfo" in deployment["status"]
+        and "extraAccessUrls" in deployment["status"]["accessInfo"]
+    ):
+        for extra_url in deployment["status"]["accessInfo"]["extraAccessUrls"]:
+            if extra_url["name"] == "in_cluster_db_access":
+                db_url = extra_url["url"].replace("http://", "")
+                return db_url
+
+    raise Exception(f"No db access endpoint found for deployment `{name}`")
+
+
+def get_db_url(app_name: str):
+    """
+    Example usage:
+        >>> from metaflow.plugins.optuna import get_db_url
+        >>> s = optuna.create_study(..., storage=get_db_url("optuna-dashboard"))
+    """
+    mf_token = auth()["x-api-key"]
+    app_url = get_deployment_db_access_endpoint(app_name)
+    return f"postgresql://userspace_default:{mf_token}@{app_url}/userspace_default?sslmode=disable"

ob_metaflow_extensions-1.4.6/metaflow_extensions/outerbounds/toplevel/plugins/optuna/__init__.py

@@ -0,0 +1 @@
+__mf_promote_submodules__ = ["plugins.optuna"]

{ob_metaflow_extensions-1.4.5 → ob_metaflow_extensions-1.4.6}/ob_metaflow_extensions.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.4.5
+Version: 1.4.6
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob_metaflow_extensions-1.4.5 → ob_metaflow_extensions-1.4.6}/ob_metaflow_extensions.egg-info/SOURCES.txt

@@ -86,6 +86,7 @@ metaflow_extensions/outerbounds/plugins/ollama/constants.py
 metaflow_extensions/outerbounds/plugins/ollama/exceptions.py
 metaflow_extensions/outerbounds/plugins/ollama/ollama.py
 metaflow_extensions/outerbounds/plugins/ollama/status_card.py
+metaflow_extensions/outerbounds/plugins/optuna/__init__.py
 metaflow_extensions/outerbounds/plugins/profilers/deco_injector.py
 metaflow_extensions/outerbounds/plugins/profilers/gpu_profile_decorator.py
 metaflow_extensions/outerbounds/plugins/profilers/simple_card_decorator.py
@@ -124,6 +125,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/ollama/__init__.py
+metaflow_extensions/outerbounds/toplevel/plugins/optuna/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/torchtune/__init__.py
 metaflow_extensions/outerbounds/toplevel/plugins/vllm/__init__.py

{ob_metaflow_extensions-1.4.5 → ob_metaflow_extensions-1.4.6}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_namespace_packages
 from pathlib import Path
 
 
-version = "1.4.5"
+version = "1.4.6"
 this_directory = Path(__file__).parent
 long_description = (this_directory / "README.md").read_text()