ob-metaflow-extensions 1.1.78__py2.py3-none-any.whl → 1.1.80__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py

@@ -0,0 +1,264 @@
+import os
+import sys
+import platform
+
+from metaflow import R, current
+from metaflow.metadata import MetaDatum
+from metaflow.metadata.util import sync_local_metadata_to_datastore
+from metaflow.sidecar import Sidecar
+from metaflow.decorators import StepDecorator
+from metaflow.exception import MetaflowException
+from metaflow.metaflow_config import (
+    DEFAULT_CONTAINER_IMAGE,
+    DEFAULT_CONTAINER_REGISTRY,
+    SNOWPARK_ACCOUNT,
+    SNOWPARK_USER,
+    SNOWPARK_PASSWORD,
+    SNOWPARK_ROLE,
+    SNOWPARK_DATABASE,
+    SNOWPARK_WAREHOUSE,
+    SNOWPARK_SCHEMA,
+)
+
+from metaflow.metaflow_config import (
+    DATASTORE_LOCAL_DIR,
+)
+
+from .snowpark_exceptions import SnowflakeException
+from metaflow.plugins.aws.aws_utils import get_docker_registry
+
+
+class Snowflake(object):
+    def __init__(self, connection_params):
+        self.connection_params = connection_params
+
+    def session(self):
+        # if using the pypi/conda decorator with @snowpark of any step,
+        # make sure to pass {'snowflake': '0.11.0'} as well to that step
+        try:
+            from snowflake.snowpark import Session
+
+            session = Session.builder.configs(self.connection_params).create()
+            return session
+        except (NameError, ImportError, ModuleNotFoundError):
+            raise SnowflakeException(
+                "Could not import module 'snowflake'.\n\nInstall Snowflake "
+                "Python package (https://pypi.org/project/snowflake/) first.\n"
+                "You can install the module by using the @pypi decorator - "
+                "Eg: @pypi(packages={'snowflake': '0.11.0'})\n"
+            )
+
+
+class SnowparkDecorator(StepDecorator):
+    name = "snowpark"
+
+    defaults = {
+        "account": None,
+        "user": None,
+        "password": None,
+        "role": None,
+        "database": None,
+        "warehouse": None,
+        "schema": None,
+        "image": None,
+        "stage": None,
+        "compute_pool": None,
+        "volume_mounts": None,
+        "external_integration": None,
+        "cpu": None,
+        "gpu": None,
+        "memory": None,
+    }
+
+    package_url = None
+    package_sha = None
+    run_time_limit = None
+
+    def __init__(self, attributes=None, statically_defined=False):
+        super(SnowparkDecorator, self).__init__(attributes, statically_defined)
+
+        if not self.attributes["account"]:
+            self.attributes["account"] = SNOWPARK_ACCOUNT
+        if not self.attributes["user"]:
+            self.attributes["user"] = SNOWPARK_USER
+        if not self.attributes["password"]:
+            self.attributes["password"] = SNOWPARK_PASSWORD
+        if not self.attributes["role"]:
+            self.attributes["role"] = SNOWPARK_ROLE
+        if not self.attributes["database"]:
+            self.attributes["database"] = SNOWPARK_DATABASE
+        if not self.attributes["warehouse"]:
+            self.attributes["warehouse"] = SNOWPARK_WAREHOUSE
+        if not self.attributes["schema"]:
+            self.attributes["schema"] = SNOWPARK_SCHEMA
+
+        # If no docker image is explicitly specified, impute a default image.
+        if not self.attributes["image"]:
+            # If metaflow-config specifies a docker image, just use that.
+            if DEFAULT_CONTAINER_IMAGE:
+                self.attributes["image"] = DEFAULT_CONTAINER_IMAGE
+            # If metaflow-config doesn't specify a docker image, assign a
+            # default docker image.
+            else:
+                # Metaflow-R has its own default docker image (rocker family)
+                if R.use_r():
+                    self.attributes["image"] = R.container_image()
+                # Default to vanilla Python image corresponding to major.minor
+                # version of the Python interpreter launching the flow.
+                self.attributes["image"] = "python:%s.%s" % (
+                    platform.python_version_tuple()[0],
+                    platform.python_version_tuple()[1],
+                )
+
+        # Assign docker registry URL for the image.
+        if not get_docker_registry(self.attributes["image"]):
+            if DEFAULT_CONTAINER_REGISTRY:
+                self.attributes["image"] = "%s/%s" % (
+                    DEFAULT_CONTAINER_REGISTRY.rstrip("/"),
+                    self.attributes["image"],
+                )
+
+    # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
+    # to understand where these functions are invoked in the lifecycle of a
+    # Metaflow flow.
+    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
+        # Set internal state.
+        self.logger = logger
+        self.environment = environment
+        self.step = step
+        self.flow_datastore = flow_datastore
+
+        if any([deco.name == "parallel" for deco in decos]):
+            raise MetaflowException(
+                "Step *{step}* contains a @parallel decorator "
+                "with the @snowpark decorator. @parallel is not supported with @snowpark.".format(
+                    step=step
+                )
+            )
+
+    def package_init(self, flow, step_name, environment):
+        try:
+            # Snowflake is a soft dependency.
+            from snowflake.snowpark import Session
+        except (NameError, ImportError, ModuleNotFoundError):
+            raise SnowflakeException(
+                "Could not import module 'snowflake'.\n\nInstall Snowflake "
+                "Python package (https://pypi.org/project/snowflake/) first.\n"
+                "You can install the module by executing - "
+                "%s -m pip install snowflake\n"
+                "or equivalent through your favorite Python package manager."
+                % sys.executable
+            )
+
+    def runtime_init(self, flow, graph, package, run_id):
+        # Set some more internal state.
+        self.flow = flow
+        self.graph = graph
+        self.package = package
+        self.run_id = run_id
+
+    def runtime_task_created(
+        self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context
+    ):
+        if not is_cloned:
+            self._save_package_once(self.flow_datastore, self.package)
+
+    def runtime_step_cli(
+        self, cli_args, retry_count, max_user_code_retries, ubf_context
+    ):
+        if retry_count <= max_user_code_retries:
+            cli_args.commands = ["snowpark", "step"]
+            cli_args.command_args.append(self.package_sha)
+            cli_args.command_args.append(self.package_url)
+            cli_args.command_options.update(self.attributes)
+            cli_args.command_options["run-time-limit"] = self.run_time_limit
+            if not R.use_r():
+                cli_args.entrypoint[0] = sys.executable
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_retries,
+        ubf_context,
+        inputs,
+    ):
+        self.metadata = metadata
+        self.task_datastore = task_datastore
+
+        # this path will exist within snowpark container services
+        login_token = open("/snowflake/session/token", "r").read()
+        connection_params = {
+            "account": os.environ.get("SNOWFLAKE_ACCOUNT"),
+            "host": os.environ.get("SNOWFLAKE_HOST"),
+            "authenticator": "oauth",
+            "token": login_token,
+            "database": os.environ.get("SNOWFLAKE_DATABASE"),
+            "schema": os.environ.get("SNOWFLAKE_SCHEMA"),
+            "autocommit": True,
+            "client_session_keep_alive": True,
+        }
+
+        # SNOWFLAKE_WAREHOUSE is injected explicitly by us
+        # but is not really required. So if it exists, we use it in
+        # connection parameters
+        if os.environ.get("SNOWFLAKE_WAREHOUSE"):
+            connection_params["warehouse"] = os.environ.get("SNOWFLAKE_WAREHOUSE")
+
+        snowflake_obj = Snowflake(connection_params)
+        current._update_env({"snowflake": snowflake_obj})
+
+        meta = {}
+        if "METAFLOW_SNOWPARK_WORKLOAD" in os.environ:
+            meta["snowflake-account"] = os.environ.get("SNOWFLAKE_ACCOUNT")
+            meta["snowflake-database"] = os.environ.get("SNOWFLAKE_DATABASE")
+            meta["snowflake-schema"] = os.environ.get("SNOWFLAKE_SCHEMA")
+            meta["snowflake-host"] = os.environ.get("SNOWFLAKE_HOST")
+            meta["snowflake-service-name"] = os.environ.get("SNOWFLAKE_SERVICE_NAME")
+
+            self._save_logs_sidecar = Sidecar("save_logs_periodically")
+            self._save_logs_sidecar.start()
+
+        if len(meta) > 0:
+            entries = [
+                MetaDatum(
+                    field=k,
+                    value=v,
+                    type=k,
+                    tags=["attempt_id:{0}".format(retry_count)],
+                )
+                for k, v in meta.items()
+                if v is not None
+            ]
+            # Register book-keeping metadata for debugging.
+            metadata.register_metadata(run_id, step_name, task_id, entries)
+
+    def task_finished(
+        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
+    ):
+        if "METAFLOW_SNOWPARK_WORKLOAD" in os.environ:
+            if hasattr(self, "metadata") and self.metadata.TYPE == "local":
+                # Note that the datastore is *always* Amazon S3 (see
+                # runtime_task_created function).
+                sync_local_metadata_to_datastore(
+                    DATASTORE_LOCAL_DIR, self.task_datastore
+                )
+
+        try:
+            self._save_logs_sidecar.terminate()
+        except:
+            # Best effort kill
+            pass
+
+    @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]

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_exceptions.py

@@ -0,0 +1,13 @@
+from metaflow.exception import MetaflowException
+
+
+class SnowflakeException(MetaflowException):
+    headline = "Snowflake error"
+
+
+class SnowparkException(MetaflowException):
+    headline = "Snowpark error"
+
+
+class SnowparkKilledException(MetaflowException):
+    headline = "Snowpark job killed"

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py

@@ -0,0 +1,235 @@
+import time
+
+from .snowpark_client import SnowparkClient
+from .snowpark_service_spec import (
+    Container,
+    Resources,
+    SnowparkServiceSpec,
+    VolumeMount,
+)
+from .snowpark_exceptions import SnowparkException
+
+mapping = str.maketrans("0123456789", "abcdefghij")
+
+
+# keep only alpha numeric characters and underscores..
+def sanitize_name(job_name: str):
+    return "".join(char for char in job_name if char.isalnum() or char == "_")
+
+
+# this is not a decorator since the exception imports need to be inside
+# and not at the top level..
+def retry_operation(
+    exception_type, func, max_retries=3, retry_delay=2, *args, **kwargs
+):
+    for attempt in range(max_retries):
+        try:
+            return func(*args, **kwargs)
+        except exception_type as e:
+            if attempt < max_retries - 1:
+                time.sleep(retry_delay)
+            else:
+                raise e
+
+
+class SnowparkJob(object):
+    def __init__(self, client: SnowparkClient, name, command, **kwargs):
+        self.client = client
+        self.name = sanitize_name(name)
+        self.command = command
+        self.kwargs = kwargs
+        self.container_name = self.name.translate(mapping).lower()
+
+    def create_job_spec(self):
+        if self.kwargs.get("image") is None:
+            raise SnowparkException(
+                "Unable to launch job on Snowpark Container Services. No docker 'image' specified."
+            )
+
+        if self.kwargs.get("stage") is None:
+            raise SnowparkException(
+                "Unable to launch job on Snowpark Container Services. No 'stage' specified."
+            )
+
+        if self.kwargs.get("compute_pool") is None:
+            raise SnowparkException(
+                "Unable to launch job on Snowpark Container Services. No 'compute_pool' specified."
+            )
+
+        resources = Resources(
+            requests={
+                k: v
+                for k, v in [
+                    ("cpu", self.kwargs.get("cpu")),
+                    ("nvidia.com/gpu", self.kwargs.get("gpu")),
+                    ("memory", self.kwargs.get("memory")),
+                ]
+                if v
+            },
+            limits={
+                k: v
+                for k, v in [
+                    ("nvidia.com/gpu", self.kwargs.get("gpu")),
+                ]
+                if v
+            },
+        )
+
+        volume_mounts = self.kwargs.get("volume_mounts")
+        vm_objs = []
+        if volume_mounts:
+            if isinstance(volume_mounts, str):
+                volume_mounts = [volume_mounts]
+            for vm in volume_mounts:
+                name, mount_path = vm.split(":", 1)
+                vm_objs.append(VolumeMount(name=name, mount_path=mount_path))
+
+        container = (
+            Container(name=self.container_name, image=self.kwargs.get("image"))
+            .env(self.kwargs.get("environment_variables"))
+            .resources(resources)
+            .volume_mounts(vm_objs)
+            .command(self.command)
+        )
+
+        self.spec = SnowparkServiceSpec().containers([container])
+        return self
+
+    def environment_variable(self, name, value):
+        # Never set to None
+        if value is None:
+            return self
+        self.kwargs["environment_variables"] = dict(
+            self.kwargs.get("environment_variables", {}), **{name: value}
+        )
+        return self
+
+    def create(self):
+        return self.create_job_spec()
+
+    def execute(self):
+        query_id, service_name = self.client.submit(
+            self.name,
+            self.spec,
+            self.kwargs.get("stage"),
+            self.kwargs.get("compute_pool"),
+            self.kwargs.get("external_integration"),
+        )
+        return RunningJob(
+            client=self.client,
+            query_id=query_id,
+            service_name=service_name,
+            **self.kwargs
+        )
+
+    def image(self, image):
+        self.kwargs["image"] = image
+        return self
+
+    def stage(self, stage):
+        self.kwargs["stage"] = stage
+        return self
+
+    def compute_pool(self, compute_pool):
+        self.kwargs["compute_pool"] = compute_pool
+        return self
+
+    def volume_mounts(self, volume_mounts):
+        self.kwargs["volume_mounts"] = volume_mounts
+        return self
+
+    def external_integration(self, external_integration):
+        self.kwargs["external_integration"] = external_integration
+        return self
+
+    def cpu(self, cpu):
+        self.kwargs["cpu"] = cpu
+        return self
+
+    def gpu(self, gpu):
+        self.kwargs["gpu"] = gpu
+        return self
+
+    def memory(self, memory):
+        self.kwargs["memory"] = memory
+        return self
+
+
+class RunningJob(object):
+    def __init__(self, client, query_id, service_name, **kwargs):
+        self.client = client
+        self.query_id = query_id
+        self.service_name = service_name
+        self.kwargs = kwargs
+
+        from snowflake.core.exceptions import NotFoundError
+
+        self.service = retry_operation(NotFoundError, self.__get_service)
+
+    def __get_service(self):
+        db = self.client.session.get_current_database()
+        schema = self.client.session.get_current_schema()
+        return (
+            self.client.root.databases[db].schemas[schema].services[self.service_name]
+        )
+
+    def __repr__(self):
+        return "{}('{}')".format(self.__class__.__name__, self.query_id)
+
+    @property
+    def id(self):
+        return self.query_id
+
+    @property
+    def job_name(self):
+        return self.service_name
+
+    def status_obj(self, timeout=0):
+        from snowflake.core.exceptions import APIError, NotFoundError
+
+        try:
+            return retry_operation(
+                APIError, self.service.get_service_status, timeout=timeout
+            )
+        except NotFoundError:
+            raise SnowparkException(
+                "The image *%s* most probably doesn't exist on Snowpark, or too many resources (CPU, GPU, memory) were requested."
+                % self.kwargs.get("image")
+            )
+
+    @property
+    def status(self):
+        return self.status_obj()[0].get("status")
+
+    @property
+    def message(self):
+        return self.status_obj()[0].get("message")
+
+    @property
+    def is_waiting(self):
+        return self.status == "PENDING"
+
+    @property
+    def is_running(self):
+        return self.status in ["PENDING", "READY"]
+
+    @property
+    def has_failed(self):
+        return self.status == "FAILED"
+
+    @property
+    def has_succeeded(self):
+        return self.status == "DONE"
+
+    @property
+    def has_finished(self):
+        return self.has_succeeded or self.has_failed
+
+    def kill(self):
+        from snowflake.core.exceptions import NotFoundError
+
+        try:
+            if not self.has_finished:
+                self.client.terminate_job(service=self.service)
+        except (NotFoundError, TypeError):
+            pass