ob-metaflow 2.12.27.1__py2.py3-none-any.whl → 2.12.30.2__py2.py3-none-any.whl

metaflow/plugins/datatools/local.py

@@ -81,6 +81,8 @@ class Local(object):
     In the future, we may want to allow it to be used in a way similar to the S3() client.
     """
 
+    TYPE = "local"
+
     @staticmethod
     def _makedirs(path):
         try:

metaflow/plugins/datatools/s3/s3.py

@@ -504,6 +504,8 @@ class S3(object):
         If `run` is not specified, use this as the S3 prefix.
     """
 
+    TYPE = "s3"
+
     @classmethod
     def get_root_from_config(cls, echo, create_on_absent=True):
         return DATATOOLS_S3ROOT

metaflow/plugins/env_escape/__init__.py

@@ -124,9 +124,9 @@ def load():
     cur_path = os.path.dirname(__file__)
     sys.path = [p for p in old_paths if p != cur_path]
     # Handle special case where we launch a shell (including with a command)
-    # and we are in the CWD (searched if '' is the first element of sys.path)
-    if cur_path == os.getcwd() and sys.path[0] == '':
-        sys.path = sys.path[1:]
+    # and we are in the CWD (searched if '' is present in sys.path)
+    if cur_path == os.getcwd() and '' in sys.path:
+        sys.path.remove("")
 
     # Remove the module (this file) to reload it properly. Do *NOT* update sys.modules but
     # modify directly since it may be referenced elsewhere

metaflow/plugins/gcp/includefile_support.py

@@ -8,6 +8,9 @@ from metaflow.exception import MetaflowException, MetaflowInternalError
 
 
 class GS(object):
+
+    TYPE = "gs"
+
     @classmethod
     def get_root_from_config(cls, echo, create_on_absent=True):
         from metaflow.metaflow_config import DATATOOLS_GSROOT

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -9,7 +9,7 @@ import metaflow.tracing as tracing
 from metaflow import JSONTypeClass, util
 from metaflow._vendor import click
 from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, MetaflowException
-from metaflow.metadata.util import sync_local_metadata_from_datastore
+from metaflow.metadata_provider.util import sync_local_metadata_from_datastore
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, KUBERNETES_LABELS
 from metaflow.mflog import TASK_LOG_SOURCE
 from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -7,8 +7,8 @@ import time
 from metaflow import current
 from metaflow.decorators import StepDecorator
 from metaflow.exception import MetaflowException
-from metaflow.metadata import MetaDatum
-from metaflow.metadata.util import sync_local_metadata_to_datastore
+from metaflow.metadata_provider import MetaDatum
+from metaflow.metadata_provider.util import sync_local_metadata_to_datastore
 from metaflow.metaflow_config import (
     DATASTORE_LOCAL_DIR,
     KUBERNETES_CONTAINER_IMAGE,
@@ -75,8 +75,9 @@ class KubernetesDecorator(StepDecorator):
         in Metaflow configuration.
     node_selector: Union[Dict[str,str], str], optional, default None
         Kubernetes node selector(s) to apply to the pod running the task.
-        Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
-        or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
+        Can be passed in as a comma separated string of values e.g.
+        'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
+        {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
     namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
         Kubernetes namespace to use when launching pod in Kubernetes.
     gpu : int, optional, default None

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -4,7 +4,6 @@ import math
 import random
 import time
 from collections import namedtuple
-
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import KUBERNETES_JOBSET_GROUP, KUBERNETES_JOBSET_VERSION
 from metaflow.tracing import inject_tracing_vars
@@ -320,33 +319,49 @@ class RunningJobSet(object):
     def kill(self):
         plural = "jobsets"
         client = self._client.get()
-        # Get the jobset
-        with client.ApiClient() as api_client:
-            api_instance = client.CustomObjectsApi(api_client)
-            try:
-                obj = api_instance.get_namespaced_custom_object(
-                    group=self._group,
-                    version=self._version,
-                    namespace=self._namespace,
-                    plural=plural,
-                    name=self._name,
-                )
-
-                # Suspend the jobset
-                obj["spec"]["suspend"] = True
-
-                api_instance.replace_namespaced_custom_object(
-                    group=self._group,
-                    version=self._version,
-                    namespace=self._namespace,
-                    plural=plural,
-                    name=obj["metadata"]["name"],
-                    body=obj,
-                )
-            except Exception as e:
-                raise KubernetesJobsetException(
-                    "Exception when suspending existing jobset: %s\n" % e
-                )
+        try:
+            # Killing the control pod will trigger the jobset to mark everything as failed.
+            # Since jobsets have a successPolicy set to `All` which ensures that everything has
+            # to succeed for the jobset to succeed.
+            from kubernetes.stream import stream
+
+            control_pod = self._fetch_pod()
+            stream(
+                client.CoreV1Api().connect_get_namespaced_pod_exec,
+                name=control_pod["metadata"]["name"],
+                namespace=control_pod["metadata"]["namespace"],
+                command=[
+                    "/bin/sh",
+                    "-c",
+                    "/sbin/killall5",
+                ],
+                stderr=True,
+                stdin=False,
+                stdout=True,
+                tty=False,
+            )
+        except Exception as e:
+            with client.ApiClient() as api_client:
+                # If we are unable to kill the control pod then
+                # Delete the jobset to kill the subsequent pods.
+                # There are a few reasons for deleting a jobset to kill it :
+                # 1. Jobset has a `suspend` attribute to suspend it's execution, but this
+                # doesn't play nicely when jobsets are deployed with other components like kueue.
+                # 2. Jobset doesn't play nicely when we mutate status
+                # 3. Deletion is a gaurenteed way of removing any pods.
+                api_instance = client.CustomObjectsApi(api_client)
+                try:
+                    api_instance.delete_namespaced_custom_object(
+                        group=self._group,
+                        version=self._version,
+                        namespace=self._namespace,
+                        plural=plural,
+                        name=self._name,
+                    )
+                except Exception as e:
+                    raise KubernetesJobsetException(
+                        "Exception when deleting existing jobset: %s\n" % e
+                    )
 
     @property
     def id(self):

metaflow/plugins/{metadata → metadata_providers}/local.py

@@ -8,9 +8,9 @@ import time
 from collections import namedtuple
 
 from metaflow.exception import MetaflowInternalError, MetaflowTaggingError
-from metaflow.metadata.metadata import ObjectOrder
+from metaflow.metadata_provider.metadata import ObjectOrder
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
-from metaflow.metadata import MetadataProvider
+from metaflow.metadata_provider import MetadataProvider
 from metaflow.tagging_util import MAX_USER_TAG_SET_SIZE, validate_tags
 
 

metaflow/plugins/{metadata → metadata_providers}/service.py

@@ -14,8 +14,8 @@ from metaflow.metaflow_config import (
     SERVICE_HEADERS,
     SERVICE_URL,
 )
-from metaflow.metadata import MetadataProvider
-from metaflow.metadata.heartbeat import HB_URL_KEY
+from metaflow.metadata_provider import MetadataProvider
+from metaflow.metadata_provider.heartbeat import HB_URL_KEY
 from metaflow.sidecar import Message, MessageTypes, Sidecar
 
 from metaflow.util import version_parse

metaflow/plugins/parallel_decorator.py

@@ -2,7 +2,7 @@ from collections import namedtuple
 from metaflow.decorators import StepDecorator
 from metaflow.unbounded_foreach import UBF_CONTROL, CONTROL_TASK_TAG
 from metaflow.exception import MetaflowException
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_current import current, Parallel
 import os
 import sys

metaflow/plugins/pypi/conda_decorator.py

@@ -8,7 +8,7 @@ import tempfile
 
 from metaflow.decorators import FlowDecorator, StepDecorator
 from metaflow.extension_support import EXT_PKG
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_environment import InvalidEnvironmentException
 from metaflow.util import get_metaflow_root
 

metaflow/plugins/test_unbounded_foreach_decorator.py

@@ -15,7 +15,7 @@ from metaflow.unbounded_foreach import (
     CONTROL_TASK_TAG,
 )
 from metaflow.util import to_unicode
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 
 
 class InternalTestUnboundedForeachInput(UnboundedForeachInput):

metaflow/runner/click_api.py

@@ -193,6 +193,10 @@ class MetaflowAPI(object):
     def chain(self):
         return self._chain
 
+    @property
+    def name(self):
+        return self._API_NAME
+
     @classmethod
     def from_cli(cls, flow_file: str, cli_collection: Callable) -> Callable:
         flow_cls = extract_flow_class_from_file(flow_file)