ob-metaflow-extensions 1.3.1__tar.gz → 1.4.34__tar.gz

{ob_metaflow_extensions-1.3.1 → ob_metaflow_extensions-1.4.34}/PKG-INFO

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

{ob_metaflow_extensions-1.3.1 → ob_metaflow_extensions-1.4.34}/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.3.1 → ob_metaflow_extensions-1.4.34}/metaflow_extensions/outerbounds/plugins/apps/core/_state_machine.py

@@ -180,6 +180,7 @@ class WorkerInfoDict(TypedDict):
     pending: Dict[str, List[WorkerStatus]]
     running: Dict[str, List[WorkerStatus]]
     crashlooping: Dict[str, List[WorkerStatus]]
+    failed: Dict[str, List[WorkerStatus]]
 
 
 class CurrentWorkerInfo(TypedDict):
@@ -442,9 +443,11 @@ def _capsule_worker_semantic_status(
                 xx[worker_version].append(w)
         return xx
 
+    # phases can be Pending, Running, Succeeded, Failed, Unknown, CrashLoopBackOff
     pending_workers = _make_version_dict(workers, "Pending")
     running_workers = _make_version_dict(workers, "Running")
     crashlooping_workers = _make_version_dict(workers, "CrashLoopBackOff")
+    failed_workers = _make_version_dict(workers, "Failed")
 
     # current_status (formulated basis):
     # - at least one pods are pending for `_end_state_capsule_version`
@@ -464,7 +467,8 @@ def _capsule_worker_semantic_status(
         "at_least_one_running": (
             count_for_version(running_workers) >= min(min_replicas, 1)
         ),
-        "at_least_one_crashlooping": count_for_version(crashlooping_workers) > 0,
+        "at_least_one_crashlooping": count_for_version(crashlooping_workers) > 0
+        or count_for_version(failed_workers) > 0,
         "none_present": (
             count_for_version(running_workers) == 0
             and count_for_version(pending_workers) == 0
@@ -484,6 +488,7 @@ def _capsule_worker_semantic_status(
             "pending": count_for_version(pending_workers),
             "running": count_for_version(running_workers),
             "crashlooping": count_for_version(crashlooping_workers),
+            "failed": count_for_version(failed_workers),
         },
     }
 
@@ -491,6 +496,7 @@ def _capsule_worker_semantic_status(
         "pending": pending_workers,
         "running": running_workers,
         "crashlooping": crashlooping_workers,
+        "failed": failed_workers,
     }
 
     return {

{ob_metaflow_extensions-1.3.1 → ob_metaflow_extensions-1.4.34}/metaflow_extensions/outerbounds/plugins/apps/core/app_cli.py

@@ -339,7 +339,7 @@ def deployment_instance_options(func):
         "--readiness-wait-time",
         type=int,
         help="The time (in seconds) to monitor the deployment for readiness after the readiness condition is met.",
-        default=4,
+        default=15,
     )
     @click.option(
         "--deployment-timeout",
@@ -377,8 +377,8 @@ def _package_necessary_things(app_config: AppConfig, logger):
     #      from the command line and that will alleviate any need to package any other directories for
     #
 
-    package_dir = app_config.get_state("packaging_directory")
-    if package_dir is None:
+    package_dirs = app_config.get_state("packaging_directories")
+    if package_dirs is None:
         app_config.set_state("code_package_url", None)
         app_config.set_state("code_package_key", None)
         return
@@ -391,13 +391,26 @@ def _package_necessary_things(app_config: AppConfig, logger):
         datastore_type=DEFAULT_DATASTORE, code_package_prefix=CODE_PACKAGE_PREFIX
     )
     package_url, package_key = packager.store(
-        paths_to_include=[package_dir], file_suffixes=suffixes
+        paths_to_include=package_dirs, file_suffixes=suffixes
     )
     app_config.set_state("code_package_url", package_url)
     app_config.set_state("code_package_key", package_key)
     logger("💾 Code package saved to : %s" % app_config.get_state("code_package_url"))
 
 
+def _sniff_pyproject_and_requirements(packaging_directories: List[str]):
+    pyproject_path = None
+    requirements_path = None
+    for directory in packaging_directories:
+        pyproject_toml = os.path.join(directory, "pyproject.toml")
+        requirements_txt = os.path.join(directory, "requirements.txt")
+        if os.path.exists(pyproject_toml):
+            pyproject_path = pyproject_toml
+        elif os.path.exists(requirements_txt):
+            requirements_path = requirements_txt
+    return pyproject_path, requirements_path
+
+
 @app.command(help="Deploy an app to the Outerbounds Platform.")
 @click.option(
     "--config-file",
@@ -449,23 +462,19 @@ def deploy(
             system_msg=True,
         )
 
-        packaging_directory = None
-        package_src_path = app_config.get("package", {}).get("src_path", None)
-        if package_src_path:
-            if os.path.isfile(package_src_path):
-                raise AppConfigError("src_path must be a directory, not a file")
-            elif os.path.isdir(package_src_path):
-                packaging_directory = os.path.abspath(package_src_path)
-            else:
-                raise AppConfigError(f"src_path '{package_src_path}' does not exist")
-        else:
-            # If src_path is None then we assume then we can assume for the moment
+        package_src_paths = app_config.get("package", {}).get("src_paths", [])
+        if package_src_paths is None:
+            package_src_paths = []
+
+        if len(package_src_paths) == 0:
+            # If src_paths is None then we assume then we can assume for the moment
             # that we can package the current working directory.
-            packaging_directory = os.getcwd()
+            package_src_paths = [os.getcwd()]
 
-        app_config.set_state("packaging_directory", packaging_directory)
+        app_config.set_state("packaging_directories", package_src_paths)
         logger(
-            "📦 Packaging directory : %s" % app_config.get_state("packaging_directory"),
+            "📦 Packaging directories : %s"
+            % ", ".join(app_config.get_state("packaging_directories")),
         )
 
         if app_config.get("no_deps", False):
@@ -484,22 +493,32 @@ def deploy(
                     dependencies.get("conda", None) is None,
                 ]
             ):
+                python_version = dependencies.get(
+                    "python"
+                )  # python gets a default value so it's always set.
                 # The user has not set any dependencies, so we can sniff the packaging directory
                 # for a dependencies file.
-                requirements_file = os.path.join(
-                    packaging_directory, "requirements.txt"
+                pyproject_toml, requirements_file = _sniff_pyproject_and_requirements(
+                    package_src_paths
                 )
-                pyproject_toml = os.path.join(packaging_directory, "pyproject.toml")
-                if os.path.exists(pyproject_toml):
+                if pyproject_toml:
                     app_config.set_state(
-                        "dependencies", {"from_pyproject_toml": pyproject_toml}
+                        "dependencies",
+                        {
+                            "from_pyproject_toml": pyproject_toml,
+                            "python": python_version,
+                        },
                     )
                     logger(
                         "📦 Using dependencies from pyproject.toml: %s" % pyproject_toml
                     )
-                elif os.path.exists(requirements_file):
+                elif requirements_file:
                     app_config.set_state(
-                        "dependencies", {"from_requirements_file": requirements_file}
+                        "dependencies",
+                        {
+                            "from_requirements_file": requirements_file,
+                            "python": python_version,
+                        },
                     )
                     logger(
                         "📦 Using dependencies from requirements.txt: %s"

{ob_metaflow_extensions-1.3.1 → ob_metaflow_extensions-1.4.34}/metaflow_extensions/outerbounds/plugins/apps/core/app_config.py

@@ -45,10 +45,11 @@ def _try_loading_yaml(file):
 class AuthType:
     BROWSER = "Browser"
     API = "API"
+    BROWSER_AND_API = "BrowserAndApi"
 
     @classmethod
     def enums(cls):
-        return [cls.BROWSER, cls.API]
+        return [cls.BROWSER, cls.API, cls.BROWSER_AND_API]
 
     @classproperty
     def default(cls):

{ob_metaflow_extensions-1.3.1 → ob_metaflow_extensions-1.4.34}/metaflow_extensions/outerbounds/plugins/apps/core/capsule.py

@@ -7,7 +7,7 @@ import sys
 import time
 from functools import partial
 import shlex
-from typing import Optional, List, Dict, Any, Tuple, Union
+from typing import Optional, List, Dict, Any, Tuple, Union, Callable
 from .utils import TODOException, safe_requests_wrapper, MaximumRetriesExceeded
 from .app_config import AppConfig, CAPSULE_DEBUG, AuthType
 from . import experimental
@@ -75,7 +75,6 @@ class CapsuleStateMachine:
         return self._status_trail
 
     def add_status(self, status: CapsuleStatus):
-        assert type(status) == dict, "TODO: Make this check somewhere else"
         self._status_trail.append({"timestamp": time.time(), "status": status})
 
     @property
@@ -116,7 +115,9 @@ class CapsuleStateMachine:
         pass
 
     def save_debug_info(self, state_dir: str):
-        debug_path = os.path.join(state_dir, f"debug_capsule_{self._capsule_id}.json")
+        debug_path = os.path.join(
+            state_dir, f"debug_capsule_sm_{self._capsule_id}.json"
+        )
         with open(debug_path, "w") as f:
             json.dump(self._status_trail, f, indent=4)
 
@@ -255,6 +256,12 @@ class CapsuleInput:
             replicas.get("min"),
             replicas.get("max"),
         )
+        rpm = replicas.get("scaling_policy", {}).get("rpm", None)
+        autoscaling_config = {}
+        if rpm:
+            autoscaling_config = {
+                "requestRateBasedAutoscalingConfig": {"targetRequestsPerMinute": rpm}
+            }
         if fixed is not None:
             _min, _max = fixed, fixed
         gpu_resource = app_config.get_state("resources").get("gpu")
@@ -296,6 +303,7 @@ class CapsuleInput:
             "autoscalingConfig": {
                 "minReplicas": _min,
                 "maxReplicas": _max,
+                **autoscaling_config,
             },
             **_scheduling_config,
             "containerStartupConfig": {
@@ -420,7 +428,7 @@ class CapsuleApi:
                 message="Capsule JSON decode failed",
             )
 
-    def get(self, capsule_id: str):
+    def get(self, capsule_id: str) -> Dict[str, Any]:
         _url = os.path.join(self._base_url, capsule_id)
         response = self._wrapped_api_caller(
             requests.get,
@@ -439,6 +447,35 @@ class CapsuleApi:
                 message="Capsule JSON decode failed",
             )
 
+    # TODO: refactor me since name *currently(9/8/25)* is unique across capsules.
+    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,
@@ -641,7 +678,7 @@ class CapsuleDeployer:
         auth_type = self._app_config.get_state("auth", {}).get("type", AuthType.default)
         if auth_type == AuthType.BROWSER:
             return "App"
-        elif auth_type == AuthType.API:
+        elif auth_type == AuthType.API or auth_type == AuthType.BROWSER_AND_API:
             return "Endpoint"
         else:
             raise TODOException(f"Unknown auth type: {auth_type}")
@@ -691,17 +728,51 @@ class CapsuleDeployer:
                 f"A capsule upgrade was triggered outside current deployment instance. Current deployment version was discarded. Current deployment version: {current_deployment_instance_version} and new version: {capsule_response.get('version', None)}",
             )
 
+    def _update_capsule_and_worker_sm(
+        self,
+        capsule_sm: "CapsuleStateMachine",
+        workers_sm: "CapsuleWorkersStateMachine",
+        logger: Callable[[str], None],
+    ) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
+        capsule_response = self.get()
+        capsule_sm.add_status(capsule_response.get("status", {}))  # type: ignore
+
+        # We need to check if someone has not upgraded the capsule under the hood and
+        # the current deployment instance is invalid.
+        self._backend_version_mismatch_check(
+            capsule_response, self.current_deployment_instance_version  # type: ignore
+        )
+        workers_response = self.get_workers()
+        capsule_sm.report_current_status(logger)
+        workers_sm.add_status(workers_response)
+        workers_sm.report_current_status(logger)
+        return capsule_response, workers_response
+
+    def _publish_capsule_debug_info(
+        self,
+        capsule_sm: "CapsuleStateMachine",
+        workers_sm: "CapsuleWorkersStateMachine",
+        capsule_response: Dict[str, Any],
+    ):
+        if CAPSULE_DEBUG and self._debug_dir:
+            capsule_sm.save_debug_info(self._debug_dir)
+            workers_sm.save_debug_info(self._debug_dir)
+            debug_path = os.path.join(
+                self._debug_dir, f"debug_capsule_{self.identifier}.json"
+            )
+            with open(debug_path, "w") as f:
+                f.write(json.dumps(capsule_response, indent=4))
+
     def _monitor_worker_readiness(
         self,
         workers_sm: "CapsuleWorkersStateMachine",
+        capsule_sm: "CapsuleStateMachine",
     ):
         """returns True if the worker is crashlooping, False otherwise"""
         logger = self._logger_fn or partial(print, file=sys.stderr)
         for i in range(self._readiness_wait_time):
             time.sleep(1)
-            workers_response = self.get_workers()
-            workers_sm.add_status(workers_response)
-            workers_sm.report_current_status(logger)
+            self._update_capsule_and_worker_sm(capsule_sm, workers_sm, logger)
             if workers_sm.is_crashlooping:
                 return True
         return False
@@ -713,7 +784,7 @@ class CapsuleDeployer:
             workers_status: List[WorkerStatus],
         ):
             for worker in workers_status:
-                if worker["phase"] == "CrashLoopBackOff":
+                if worker["phase"] == "CrashLoopBackOff" or worker["phase"] == "Failed":
                     return worker["workerId"]
             return None
 
@@ -747,21 +818,19 @@ class CapsuleDeployer:
             minimum_replicas=min_replicas,
         )
         self.status = state_machine
+
+        # This loop will check all the conditions that help verify the terminal state.
+        # How it works is by extracting the statuses of the capsule and workers and
+        # then adding them as a part of a state-machine that helps track transitions and
+        # helps derive terminal states.
+        # We will first keep checking for terminal conditions or outright failure conditions
+        # If we reach a teminal condition like described in `DEPLOYMENT_READY_CONDITIONS`, then
+        # we will further check for readiness conditions.
         for i in range(self._create_timeout):
             time.sleep(1)
-            capsule_response = self.get()
-            workers_response = self.get_workers()
-
-            # We first need to check if someone has not upgraded the capsule under the hood and
-            # the current deployment instance is invalid.
-            self._backend_version_mismatch_check(
-                capsule_response, self.current_deployment_instance_version  # type: ignore
+            capsule_response, _ = self._update_capsule_and_worker_sm(
+                state_machine, workers_state_machine, logger
             )
-            state_machine.add_status(capsule_response.get("status", {}))  # type: ignore
-            workers_state_machine.add_status(workers_response)
-            state_machine.report_current_status(logger)
-
-            workers_state_machine.report_current_status(logger)
             # Deployment readiness checks will determine what is the terminal state
             # of the workerstate machine. If we detect a terminal state in the workers,
             # then even if the capsule upgrade is still in progress we will end up crashing
@@ -786,9 +855,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
@@ -802,7 +873,8 @@ class CapsuleDeployer:
                         % self.identifier
                     )
                     _further_readiness_check_failed = self._monitor_worker_readiness(
-                        workers_state_machine
+                        workers_state_machine,
+                        state_machine,
                     )
 
                 if CAPSULE_DEBUG:
@@ -846,13 +918,18 @@ class CapsuleDeployer:
 
                 break
 
-            if CAPSULE_DEBUG and self._debug_dir:
-                state_machine.save_debug_info(self._debug_dir)
-                workers_state_machine.save_debug_info(self._debug_dir)
-                if i % 3 == 0:  # Every 3 seconds report the status
-                    logger(
-                        f"[debug] 💊 {self.capsule_type} {self.identifier} deployment status: {state_machine.current_status} | worker states: {workers_state_machine.current_status}"
-                    )
+            self._publish_capsule_debug_info(
+                state_machine, workers_state_machine, capsule_response
+            )
+
+            if CAPSULE_DEBUG and i % 3 == 0:  # Every 3 seconds report the status
+                logger(
+                    f"[debug] 💊 {self.capsule_type} {self.identifier} deployment status: {state_machine.current_status} | worker states: {workers_state_machine.current_status} | capsule_ready : {capsule_ready} | further_check_worker_readiness {further_check_worker_readiness}"
+                )
+
+        self._publish_capsule_debug_info(
+            state_machine, workers_state_machine, capsule_response
+        )
 
         # We will only check ready_to_serve_traffic under the following conditions:
         # If the readiness condition is not Async and min_replicas in this deployment
@@ -870,13 +947,6 @@ class CapsuleDeployer:
                 f"Capsule {self.identifier} failed to be ready to serve traffic",
             )
 
-        if CAPSULE_DEBUG and self._debug_dir:
-            state_machine.save_debug_info(self._debug_dir)
-            workers_state_machine.save_debug_info(self._debug_dir)
-            logger(
-                f"[debug] 💊 {self.capsule_type} {self.identifier} deployment status [on return]: {state_machine.current_status} | worker states: {workers_state_machine.current_status}"
-            )
-
         return dict(
             id=self.identifier,
             auth_type=self.capsule_type,

{ob_metaflow_extensions-1.3.1 → ob_metaflow_extensions-1.4.34}/metaflow_extensions/outerbounds/plugins/apps/core/code_package/code_packager.py

@@ -381,6 +381,7 @@ class CodePackager:
                     "google-auth",
                     "simple-gcp-object-downloader",
                     "google-cloud-secret-manager",
+                    "packaging",
                 ],
                 "local": [],
             }
@@ -441,6 +442,7 @@ class CodePackager:
         root,
         exclude_hidden=True,
         suffixes=None,
+        normalized_rel_path=False,
     ) -> List[Tuple[str, str]]:
         """
         Walk a directory and yield tuples of (file_path, relative_arcname) for files
@@ -456,6 +458,8 @@ class CodePackager:
             Whether to exclude hidden files and directories (those starting with '.')
         suffixes : List[str], optional
             List of file suffixes to include (e.g. ['.py', '.txt'])
+        normalized_rel_path : bool, default False
+            Whether to normalize the relative from the root. ie if the root is /a/b/c and the file is /a/b/c/d/e.py then the relative path will be d/e.py
 
         Returns
         -------
@@ -468,6 +472,8 @@ class CodePackager:
         for file_path, rel_path in symlink_friendly_walk(
             root, exclude_hidden, suffixes
         ):
+            if normalized_rel_path:
+                rel_path = file_path.replace(root, "")
             files.append((file_path, rel_path))
         return files
 
@@ -520,6 +526,7 @@ class CodePackager:
                         path,
                         exclude_hidden=True,
                         suffixes=suffixes,
+                        normalized_rel_path=True,
                     ):
                         tar.add(
                             file_path,

{ob_metaflow_extensions-1.3.1 → ob_metaflow_extensions-1.4.34}/metaflow_extensions/outerbounds/plugins/apps/core/config/schema_export.py

@@ -276,6 +276,9 @@ def _get_field_schema(field_info) -> Dict[str, Any]:
     if field_info.behavior:
         field_schema["mutation_behavior"] = field_info.behavior
 
+    if field_info.cli_meta is not None:
+        field_schema["cli_option"] = field_info.cli_meta.cli_option_str
+
     # Handle validation from CLI metadata
     if field_info.cli_meta and field_info.validation_fn:
         # Add validation hints in description

{ob_metaflow_extensions-1.3.1 → ob_metaflow_extensions-1.4.34}/metaflow_extensions/outerbounds/plugins/apps/core/config/typed_configs.py

@@ -40,6 +40,7 @@ class ResourceConfigDict(TypedDict, total=False):
     memory: Optional[str]
     gpu: Optional[str]
     disk: Optional[str]
+    shared_memory: Optional[str]
 
 
 class AuthConfigDict(TypedDict, total=False):
@@ -51,6 +52,11 @@ class ReplicaConfigDict(TypedDict, total=False):
     fixed: Optional[int]
     min: Optional[int]
     max: Optional[int]
+    scaling_policy: Optional["ScalingPolicyConfigDict"]
+
+
+class ScalingPolicyConfigDict(TypedDict, total=False):
+    rpm: Optional[int]
 
 
 class DependencyConfigDict(TypedDict, total=False):
@@ -62,7 +68,7 @@ class DependencyConfigDict(TypedDict, total=False):
 
 
 class PackageConfigDict(TypedDict, total=False):
-    src_path: Optional[str]
+    src_paths: Optional[list]
     suffixes: Optional[list]
 
 
@@ -91,6 +97,7 @@ class TypedCoreConfig:
         branch: Optional[str] = None,
         models: Optional[list] = None,
         data: Optional[list] = None,
+        generate_static_url: Optional[bool] = None,
         **kwargs
     ) -> None:
         self._kwargs = {
@@ -116,6 +123,7 @@ class TypedCoreConfig:
             "branch": branch,
             "models": models,
             "data": data,
+            "generate_static_url": generate_static_url,
         }
         # Add any additional kwargs
         self._kwargs.update(kwargs)