ob-metaflow 2.15.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/kubernetes/kubernetes.py

@@ -90,6 +90,7 @@ class Kubernetes(object):
         step_name,
         task_id,
         attempt,
+        code_package_metadata,
         code_package_url,
         step_cmds,
     ):
@@ -104,7 +105,7 @@ class Kubernetes(object):
             stderr_path=STDERR_PATH,
         )
         init_cmds = self._environment.get_package_commands(
-            code_package_url, self._datastore.TYPE
+            code_package_url, self._datastore.TYPE, code_package_metadata
         )
         init_expr = " && ".join(init_cmds)
         step_expr = bash_capture_logs(
@@ -165,11 +166,13 @@ class Kubernetes(object):
         task_id,
         attempt,
         user,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
         docker_image,
         docker_image_pull_policy,
+        image_pull_secrets=None,
         step_cli=None,
         service_account=None,
         secrets=None,
@@ -206,6 +209,7 @@ class Kubernetes(object):
                 node_selector=node_selector,
                 image=docker_image,
                 image_pull_policy=docker_image_pull_policy,
+                image_pull_secrets=image_pull_secrets,
                 cpu=cpu,
                 memory=memory,
                 disk=disk,
@@ -230,6 +234,7 @@ class Kubernetes(object):
                 qos=qos,
                 security_context=security_context,
             )
+            .environment_variable("METAFLOW_CODE_METADATA", code_package_metadata)
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
             .environment_variable("METAFLOW_CODE_DS", code_package_ds)
@@ -429,6 +434,7 @@ class Kubernetes(object):
             step_name=step_name,
             task_id=_tskid,
             attempt=attempt,
+            code_package_metadata=code_package_metadata,
             code_package_url=code_package_url,
             step_cmds=[
                 step_cli.replace(
@@ -477,12 +483,14 @@ class Kubernetes(object):
         task_id,
         attempt,
         user,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         code_package_ds,
         step_cli,
         docker_image,
         docker_image_pull_policy,
+        image_pull_secrets=None,
         service_account=None,
         secrets=None,
         node_selector=None,
@@ -524,11 +532,13 @@ class Kubernetes(object):
                     step_name=step_name,
                     task_id=task_id,
                     attempt=attempt,
+                    code_package_metadata=code_package_metadata,
                     code_package_url=code_package_url,
                     step_cmds=[step_cli],
                 ),
                 image=docker_image,
                 image_pull_policy=docker_image_pull_policy,
+                image_pull_secrets=image_pull_secrets,
                 cpu=cpu,
                 memory=memory,
                 disk=disk,
@@ -551,6 +561,7 @@ class Kubernetes(object):
                 qos=qos,
                 security_context=security_context,
             )
+            .environment_variable("METAFLOW_CODE_METADATA", code_package_metadata)
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
             .environment_variable("METAFLOW_CODE_DS", code_package_ds)

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -41,6 +41,7 @@ def kubernetes():
 )
 @tracing.cli("kubernetes/step")
 @click.argument("step-name")
+@click.argument("code-package-metadata")
 @click.argument("code-package-sha")
 @click.argument("code-package-url")
 @click.option(
@@ -53,6 +54,12 @@ def kubernetes():
     default=None,
     help="Optional Docker Image Pull Policy for Kubernetes pod.",
 )
+@click.option(
+    "--image-pull-secrets",
+    default=None,
+    type=JSONTypeClass(),
+    multiple=False,
+)
 @click.option(
     "--service-account",
     help="IRSA requirement for Kubernetes pod.",
@@ -155,11 +162,13 @@ def kubernetes():
 def step(
     ctx,
     step_name,
+    code_package_metadata,
     code_package_sha,
     code_package_url,
     executable=None,
     image=None,
     image_pull_policy=None,
+    image_pull_secrets=None,
     service_account=None,
     secrets=None,
     node_selector=None,
@@ -297,12 +306,14 @@ def step(
                 task_id=task_id,
                 attempt=str(retry_count),
                 user=util.get_username(),
+                code_package_metadata=code_package_metadata,
                 code_package_sha=code_package_sha,
                 code_package_url=code_package_url,
                 code_package_ds=ctx.obj.flow_datastore.TYPE,
                 step_cli=step_cli,
                 docker_image=image,
                 docker_image_pull_policy=image_pull_policy,
+                image_pull_secrets=image_pull_secrets,
                 service_account=service_account,
                 secrets=secrets,
                 node_selector=node_selector,

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -11,6 +11,7 @@ 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,
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
     KUBERNETES_CONTAINER_IMAGE,
     KUBERNETES_CONTAINER_REGISTRY,
     KUBERNETES_CPU,
@@ -18,6 +19,7 @@ from metaflow.metaflow_config import (
     KUBERNETES_FETCH_EC2_METADATA,
     KUBERNETES_GPU_VENDOR,
     KUBERNETES_IMAGE_PULL_POLICY,
+    KUBERNETES_IMAGE_PULL_SECRETS,
     KUBERNETES_MEMORY,
     KUBERNETES_LABELS,
     KUBERNETES_ANNOTATIONS,
@@ -74,6 +76,10 @@ class KubernetesDecorator(StepDecorator):
         not, a default Docker image mapping to the current version of Python is used.
     image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
         If given, the imagePullPolicy to be applied to the Docker image of the step.
+    image_pull_secrets: List[str], default []
+        The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
+        Kubernetes image pull secrets to use when pulling container images
+        in Kubernetes.
     service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
         Kubernetes service account to use when launching pod in Kubernetes.
     secrets : List[str], optional, default None
@@ -92,7 +98,7 @@ class KubernetesDecorator(StepDecorator):
         the scheduled node should not have GPUs.
     gpu_vendor : str, default KUBERNETES_GPU_VENDOR
         The vendor of the GPUs to be used for this step.
-    tolerations : List[str], default []
+    tolerations : List[Dict[str,str]], default []
         The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
         Kubernetes tolerations to use when launching pod in Kubernetes.
     labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
@@ -141,6 +147,7 @@ class KubernetesDecorator(StepDecorator):
         "disk": "10240",
         "image": None,
         "image_pull_policy": None,
+        "image_pull_secrets": None,  # e.g., ["regcred"]
         "service_account": None,
         "secrets": None,  # e.g., mysecret
         "node_selector": None,  # e.g., kubernetes.io/os=linux
@@ -164,6 +171,7 @@ class KubernetesDecorator(StepDecorator):
         "qos": KUBERNETES_QOS,
         "security_context": None,
     }
+    package_metadata = None
     package_url = None
     package_sha = None
     run_time_limit = None
@@ -173,8 +181,6 @@ class KubernetesDecorator(StepDecorator):
     target_platform = KUBERNETES_CONDA_ARCH or "linux-64"
 
     def init(self):
-        super(KubernetesDecorator, self).init()
-
         if not self.attributes["namespace"]:
             self.attributes["namespace"] = KUBERNETES_NAMESPACE
         if not self.attributes["service_account"]:
@@ -194,6 +200,10 @@ class KubernetesDecorator(StepDecorator):
             )
         if not self.attributes["image_pull_policy"] and KUBERNETES_IMAGE_PULL_POLICY:
             self.attributes["image_pull_policy"] = KUBERNETES_IMAGE_PULL_POLICY
+        if not self.attributes["image_pull_secrets"] and KUBERNETES_IMAGE_PULL_SECRETS:
+            self.attributes["image_pull_secrets"] = json.loads(
+                KUBERNETES_IMAGE_PULL_SECRETS
+            )
 
         if isinstance(self.attributes["node_selector"], str):
             self.attributes["node_selector"] = parse_kube_keyvalue_list(
@@ -476,6 +486,7 @@ class KubernetesDecorator(StepDecorator):
             # to execute on Kubernetes anymore. We can execute possible fallback
             # code locally.
             cli_args.commands = ["kubernetes", "step"]
+            cli_args.command_args.append(self.package_metadata)
             cli_args.command_args.append(self.package_sha)
             cli_args.command_args.append(self.package_url)
 
@@ -494,6 +505,7 @@ class KubernetesDecorator(StepDecorator):
                         for key, val in v.items()
                     ]
                 elif k in [
+                    "image_pull_secrets",
                     "tolerations",
                     "persistent_volume_claims",
                     "labels",
@@ -646,9 +658,16 @@ class KubernetesDecorator(StepDecorator):
     @classmethod
     def _save_package_once(cls, flow_datastore, package):
         if cls.package_url is None:
-            cls.package_url, cls.package_sha = flow_datastore.save_data(
-                [package.blob], len_hint=1
-            )[0]
+            if not FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+                cls.package_url, cls.package_sha = flow_datastore.save_data(
+                    [package.blob], len_hint=1
+                )[0]
+                cls.package_metadata = package.package_metadata
+            else:
+                # Blocks until the package is uploaded
+                cls.package_url = package.package_url()
+                cls.package_sha = package.package_sha()
+                cls.package_metadata = package.package_metadata
 
 
 # TODO: Unify this method with the multi-node setup in @batch

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -235,8 +235,10 @@ class KubernetesJob(object):
                         )
                     ],
                     node_selector=self._kwargs.get("node_selector"),
-                    # TODO (savin): Support image_pull_secrets
-                    # image_pull_secrets=?,
+                    image_pull_secrets=[
+                        client.V1LocalObjectReference(secret)
+                        for secret in self._kwargs.get("image_pull_secrets") or []
+                    ],
                     # TODO (savin): Support preemption policies
                     # preemption_policy=?,
                     #
@@ -520,12 +522,10 @@ class RunningJob(object):
         # 3. If the pod object hasn't shown up yet, we set the parallelism to 0
         #    to preempt it.
         client = self._client.get()
-
         if not self.is_done:
             if self.is_running:
                 # Case 1.
                 from kubernetes.stream import stream
-
                 api_instance = client.CoreV1Api
                 try:
                     # TODO: stream opens a web-socket connection. It may
@@ -591,6 +591,10 @@ class RunningJob(object):
             return self.id
         return "job %s" % self._name
 
+    @property
+    def is_unschedulable(self):
+        return self._job["metadata"]["annotations"].get("metaflow/job_status", "") == "Unsatisfiable_Resource_Request"
+
     @property
     def is_done(self):
         # Check if the container is done. As a side effect, also refreshes self._job and
@@ -604,6 +608,7 @@ class RunningJob(object):
                 or bool(self._job["status"].get("failed"))
                 or self._are_pod_containers_done
                 or (self._job["spec"]["parallelism"] == 0)
+                or self.is_unschedulable
             )
 
         if not done():
@@ -661,6 +666,7 @@ class RunningJob(object):
             bool(self._job["status"].get("failed"))
             or self._has_any_container_failed
             or (self._job["spec"]["parallelism"] == 0)
+            or self.is_unschedulable
         )
         return retval
 
@@ -758,6 +764,8 @@ class RunningJob(object):
                 return 0, None
             # Best effort since Pod object can disappear on us at anytime
             else:
+                if self.is_unschedulable:
+                    return 1, self._job["metadata"]["annotations"].get("metaflow/job_status_reason", "")
                 if self._pod.get("status", {}).get("phase") not in (
                     "Succeeded",
                     "Failed",

metaflow/plugins/kubernetes/kubernetes_jobsets.py

@@ -6,6 +6,7 @@ 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
+from metaflow._vendor import yaml
 
 from .kube_utils import qos_requests_and_limits
 
@@ -718,8 +719,11 @@ class JobSetSpec(object):
                                     )
                                 ],
                                 node_selector=self._kwargs.get("node_selector"),
-                                # TODO (savin): Support image_pull_secrets
-                                # image_pull_secrets=?,
+                                image_pull_secrets=[
+                                    client.V1LocalObjectReference(secret)
+                                    for secret in self._kwargs.get("image_pull_secrets")
+                                    or []
+                                ],
                                 # TODO (savin): Support preemption policies
                                 # preemption_policy=?,
                                 #
@@ -1022,34 +1026,32 @@ class KubernetesArgoJobSet(object):
 
     def dump(self):
         client = self._kubernetes_sdk
-
-        data = json.dumps(
-            client.ApiClient().sanitize_for_serialization(
-                dict(
-                    apiVersion=self._group + "/" + self._version,
-                    kind="JobSet",
-                    metadata=client.api_client.ApiClient().sanitize_for_serialization(
-                        client.V1ObjectMeta(
-                            name=self.name,
-                            labels=self._labels,
-                            annotations=self._annotations,
-                        )
-                    ),
-                    spec=dict(
-                        replicatedJobs=[self.control.dump(), self.worker.dump()],
-                        suspend=False,
-                        startupPolicy=None,
-                        successPolicy=None,
-                        # The Failure Policy helps setting the number of retries for the jobset.
-                        # but we don't rely on it and instead rely on either the local scheduler
-                        # or the Argo Workflows to handle retries.
-                        failurePolicy=None,
-                        network=None,
-                    ),
-                    status=None,
-                )
+        js_dict = client.ApiClient().sanitize_for_serialization(
+            dict(
+                apiVersion=self._group + "/" + self._version,
+                kind="JobSet",
+                metadata=client.api_client.ApiClient().sanitize_for_serialization(
+                    client.V1ObjectMeta(
+                        name=self.name,
+                        labels=self._labels,
+                        annotations=self._annotations,
+                    )
+                ),
+                spec=dict(
+                    replicatedJobs=[self.control.dump(), self.worker.dump()],
+                    suspend=False,
+                    startupPolicy=None,
+                    successPolicy=None,
+                    # The Failure Policy helps setting the number of retries for the jobset.
+                    # but we don't rely on it and instead rely on either the local scheduler
+                    # or the Argo Workflows to handle retries.
+                    failurePolicy=None,
+                    network=None,
+                ),
+                status=None,
             )
         )
+        data = yaml.dump(js_dict, default_flow_style=False, indent=2)
         # The values we populate in the Jobset manifest (for Argo Workflows) piggybacks on the Argo Workflow's templating engine.
         # Even though Argo Workflows's templating helps us constructing all the necessary IDs and populating the fields
         # required by Metaflow, we run into one glitch. When we construct JSON/YAML serializable objects,
@@ -1064,7 +1066,6 @@ class KubernetesArgoJobSet(object):
         # Since the value of `num_parallel` can be dynamic and can change from run to run, we need to ensure that the
         # value can be passed-down dynamically and is **explicitly set as a integer** in the Jobset Manifest submitted as a
         # part of the Argo Workflow
-
-        quoted_substring = '"{{=asInt(inputs.parameters.workerCount)}}"'
+        quoted_substring = "'{{=asInt(inputs.parameters.workerCount)}}'"
         unquoted_substring = "{{=asInt(inputs.parameters.workerCount)}}"
         return data.replace(quoted_substring, unquoted_substring)