ob-metaflow-extensions 1.1.45rc3__py2.py3-none-any.whl → 1.5.1__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py

@@ -0,0 +1,277 @@
+import os
+import sys
+import time
+import traceback
+
+from metaflow._vendor import click
+from metaflow import util
+from metaflow import R
+from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY
+from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
+from metaflow.metadata_provider.util import sync_local_metadata_from_datastore
+from metaflow.mflog import TASK_LOG_SOURCE
+
+from .snowpark_exceptions import SnowparkKilledException
+from .snowpark import Snowpark
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.group(help="Commands related to Snowpark.")
+def snowpark():
+    pass
+
+
+@snowpark.command(
+    help="Execute a single task using Snowpark. This command calls the "
+    "top-level step command inside a Snowpark job with the given options. "
+    "Typically you do not call this command directly; it is used internally by "
+    "Metaflow."
+)
+@click.argument("step-name")
+@click.argument("code-package-sha")
+@click.argument("code-package-url")
+@click.option(
+    "--executable", help="Executable requirement for Snowpark Container Services."
+)
+# below are snowpark specific options
+@click.option(
+    "--account",
+    help="Account ID for Snowpark Container Services.",
+)
+@click.option(
+    "--user",
+    help="Username for Snowpark Container Services.",
+)
+@click.option(
+    "--password",
+    help="Password for Snowpark Container Services.",
+)
+@click.option(
+    "--role",
+    help="Role for Snowpark Container Services.",
+)
+@click.option(
+    "--database",
+    help="Database for Snowpark Container Services.",
+)
+@click.option(
+    "--warehouse",
+    help="Warehouse for Snowpark Container Services.",
+)
+@click.option(
+    "--schema",
+    help="Schema for Snowpark Container Services.",
+)
+@click.option(
+    "--integration",
+    help="Outerbounds OAuth integration name for Snowpark Container Services. When set, uses OAuth authentication instead of password.",
+)
+@click.option(
+    "--image",
+    help="Docker image requirement for Snowpark Container Services. In name:version format.",
+)
+@click.option("--stage", help="Stage requirement for Snowpark Container Services.")
+@click.option(
+    "--compute-pool", help="Compute Pool requirement for Snowpark Container Services."
+)
+@click.option("--volume-mounts", multiple=True)
+@click.option(
+    "--external-integration",
+    multiple=True,
+    help="External Integration requirement for Snowpark Container Services.",
+)
+@click.option("--cpu", help="CPU requirement for Snowpark Container Services.")
+@click.option("--gpu", help="GPU requirement for Snowpark Container Services.")
+@click.option("--memory", help="Memory requirement for Snowpark Container Services.")
+# TODO: secrets, volumes for snowpark
+# TODO: others to consider: ubf-context, num-parallel
+@click.option("--run-id", help="Passed to the top-level 'step'.")
+@click.option("--task-id", help="Passed to the top-level 'step'.")
+@click.option("--input-paths", help="Passed to the top-level 'step'.")
+@click.option("--split-index", help="Passed to the top-level 'step'.")
+@click.option("--clone-path", help="Passed to the top-level 'step'.")
+@click.option("--clone-run-id", help="Passed to the top-level 'step'.")
+@click.option(
+    "--tag", multiple=True, default=None, help="Passed to the top-level 'step'."
+)
+@click.option("--namespace", default=None, help="Passed to the top-level 'step'.")
+@click.option("--retry-count", default=0, help="Passed to the top-level 'step'.")
+@click.option(
+    "--max-user-code-retries", default=0, help="Passed to the top-level 'step'."
+)
+# TODO: this is not used anywhere as of now...
+@click.option(
+    "--run-time-limit",
+    default=5 * 24 * 60 * 60,  # Default is set to 5 days
+    help="Run time limit in seconds for Snowpark container.",
+)
+@click.pass_context
+def step(
+    ctx,
+    step_name,
+    code_package_sha,
+    code_package_url,
+    executable=None,
+    account=None,
+    user=None,
+    password=None,
+    role=None,
+    database=None,
+    warehouse=None,
+    schema=None,
+    integration=None,
+    image=None,
+    stage=None,
+    compute_pool=None,
+    volume_mounts=None,
+    external_integration=None,
+    cpu=None,
+    gpu=None,
+    memory=None,
+    run_time_limit=None,
+    **kwargs
+):
+    def echo(msg, stream="stderr", job_id=None, **kwargs):
+        msg = util.to_unicode(msg)
+        if job_id:
+            msg = "[%s] %s" % (job_id, msg)
+        ctx.obj.echo_always(msg, err=(stream == sys.stderr), **kwargs)
+
+    if R.use_r():
+        entrypoint = R.entrypoint()
+    else:
+        executable = ctx.obj.environment.executable(step_name, executable)
+        entrypoint = "%s -u %s" % (executable, os.path.basename(sys.argv[0]))
+
+    top_args = " ".join(util.dict_to_cli_options(ctx.parent.parent.params))
+
+    input_paths = kwargs.get("input_paths")
+    split_vars = None
+    if input_paths:
+        max_size = 30 * 1024
+        split_vars = {
+            "METAFLOW_INPUT_PATHS_%d" % (i // max_size): input_paths[i : i + max_size]
+            for i in range(0, len(input_paths), max_size)
+        }
+        kwargs["input_paths"] = "".join("${%s}" % s for s in split_vars.keys())
+
+    step_args = " ".join(util.dict_to_cli_options(kwargs))
+
+    step_cli = "{entrypoint} {top_args} step {step} {step_args}".format(
+        entrypoint=entrypoint,
+        top_args=top_args,
+        step=step_name,
+        step_args=step_args,
+    )
+    node = ctx.obj.graph[step_name]
+
+    # Get retry information
+    retry_count = kwargs.get("retry_count", 0)
+    retry_deco = [deco for deco in node.decorators if deco.name == "retry"]
+    minutes_between_retries = None
+    if retry_deco:
+        minutes_between_retries = int(
+            retry_deco[0].attributes.get("minutes_between_retries", 1)
+        )
+    if retry_count:
+        ctx.obj.echo_always(
+            "Sleeping %d minutes before the next retry" % minutes_between_retries
+        )
+        time.sleep(minutes_between_retries * 60)
+
+    # Set batch attributes
+    task_spec = {
+        "flow_name": ctx.obj.flow.name,
+        "step_name": step_name,
+        "run_id": kwargs["run_id"],
+        "task_id": kwargs["task_id"],
+        "retry_count": str(retry_count),
+    }
+    attrs = {"metaflow.%s" % k: v for k, v in task_spec.items()}
+    attrs["metaflow.user"] = util.get_username()
+    attrs["metaflow.version"] = ctx.obj.environment.get_environment_info()[
+        "metaflow_version"
+    ]
+
+    # Set environment
+    env = {}
+    env_deco = [deco for deco in node.decorators if deco.name == "environment"]
+    if env_deco:
+        env = env_deco[0].attributes["vars"]
+
+    # Add the environment variables related to the input-paths argument
+    if split_vars:
+        env.update(split_vars)
+
+    # Set log tailing.
+    ds = ctx.obj.flow_datastore.get_task_datastore(
+        mode="w",
+        run_id=kwargs["run_id"],
+        step_name=step_name,
+        task_id=kwargs["task_id"],
+        attempt=int(retry_count),
+    )
+    stdout_location = ds.get_log_location(TASK_LOG_SOURCE, "stdout")
+    stderr_location = ds.get_log_location(TASK_LOG_SOURCE, "stderr")
+
+    def _sync_metadata():
+        if ctx.obj.metadata.TYPE == "local":
+            sync_local_metadata_from_datastore(
+                DATASTORE_LOCAL_DIR,
+                ctx.obj.flow_datastore.get_task_datastore(
+                    kwargs["run_id"], step_name, kwargs["task_id"]
+                ),
+            )
+
+    try:
+        snowpark = Snowpark(
+            datastore=ctx.obj.flow_datastore,
+            metadata=ctx.obj.metadata,
+            environment=ctx.obj.environment,
+            client_credentials={
+                "account": account,
+                "user": user,
+                "password": password,
+                "role": role,
+                "database": database,
+                "warehouse": warehouse,
+                "schema": schema,
+                "integration": integration,
+            },
+        )
+        with ctx.obj.monitor.measure("metaflow.snowpark.launch_job"):
+            snowpark.launch_job(
+                step_name=step_name,
+                step_cli=step_cli,
+                task_spec=task_spec,
+                code_package_sha=code_package_sha,
+                code_package_url=code_package_url,
+                code_package_ds=ctx.obj.flow_datastore.TYPE,
+                image=image,
+                stage=stage,
+                compute_pool=compute_pool,
+                volume_mounts=volume_mounts,
+                external_integration=external_integration,
+                cpu=cpu,
+                gpu=gpu,
+                memory=memory,
+                run_time_limit=run_time_limit,
+                env=env,
+                attrs=attrs,
+            )
+    except Exception:
+        traceback.print_exc(chain=False)
+        _sync_metadata()
+        sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
+    try:
+        snowpark.wait(stdout_location, stderr_location, echo=echo)
+    except SnowparkKilledException:
+        # don't retry killed tasks
+        traceback.print_exc()
+        sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
+    finally:
+        _sync_metadata()

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py

@@ -0,0 +1,150 @@
+import sys
+import tempfile
+
+from .snowpark_exceptions import SnowflakeException
+from .snowpark_service_spec import generate_spec_file
+
+from metaflow.exception import MetaflowException
+
+
+class SnowparkClient(object):
+    def __init__(
+        self,
+        account: str = None,
+        user: str = None,
+        password: str = None,
+        role: str = None,
+        database: str = None,
+        warehouse: str = None,
+        schema: str = None,
+        autocommit: bool = True,
+        integration: str = None,
+    ):
+        try:
+            from snowflake.core import Root
+            from snowflake.snowpark import Session
+
+            from snowflake.connector.errors import DatabaseError
+        except (NameError, ImportError, ModuleNotFoundError):
+            raise SnowflakeException(
+                "Could not import module 'snowflake'.\n\nInstall Snowflake "
+                "Python packages first:\n"
+                "  snowflake==1.8.0\n"
+                "  snowflake-connector-python==3.18.0\n"
+                "  snowflake-snowpark-python==1.40.0\n\n"
+                "You can install them by executing:\n"
+                "%s -m pip install snowflake==1.8.0 snowflake-connector-python==3.18.0 snowflake-snowpark-python==1.40.0\n"
+                "or equivalent through your favorite Python package manager."
+                % sys.executable
+            )
+
+        if integration:
+            # Use OAuth authentication via Outerbounds integration
+            from metaflow_extensions.outerbounds.plugins.snowflake.snowflake import (
+                get_oauth_connection_params,
+            )
+
+            self.connection_parameters = get_oauth_connection_params(
+                user=user or "",
+                role=role or "",
+                integration=integration,
+                schema=schema or "",
+                account=account,
+                warehouse=warehouse,
+                database=database,
+            )
+            self.connection_parameters["autocommit"] = autocommit
+        else:
+            # Password-based authentication
+            self.connection_parameters = {
+                "account": account,
+                "user": user,
+                "password": password,
+                "role": role,
+                "warehouse": warehouse,
+                "database": database,
+                "schema": schema,
+                "autocommit": autocommit,
+            }
+
+        # Remove None values from connection parameters
+        self.connection_parameters = {
+            k: v for k, v in self.connection_parameters.items() if v is not None
+        }
+
+        try:
+            self.session = Session.builder.configs(self.connection_parameters).create()
+            self.root = Root(self.session)
+        except DatabaseError as e:
+            raise ConnectionError(e)
+
+    def __del__(self):
+        if hasattr(self, "session"):
+            self.session.close()
+
+    def __check_existence_of_stage_and_compute_pool(
+        self, db, schema, stage, compute_pool
+    ):
+        from snowflake.core.exceptions import NotFoundError
+
+        # check if stage exists, will raise an error otherwise
+        try:
+            self.root.databases[db].schemas[schema].stages[stage].fetch()
+        except NotFoundError:
+            raise MetaflowException(
+                "Stage *%s* does not exist or not authorized." % stage
+            )
+
+        # check if compute_pool exists, will raise an error otherwise
+        try:
+            self.root.compute_pools[compute_pool].fetch()
+        except NotFoundError:
+            raise MetaflowException(
+                "Compute pool *%s* does not exist or not authorized." % compute_pool
+            )
+
+    def submit(self, name: str, spec, stage, compute_pool, external_integration):
+        db = self.session.get_current_database()
+        schema = self.session.get_current_schema()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            snowpark_spec_file = tempfile.NamedTemporaryFile(
+                dir=temp_dir, delete=False, suffix=".yaml"
+            )
+            generate_spec_file(spec, snowpark_spec_file.name, format="yaml")
+
+            self.__check_existence_of_stage_and_compute_pool(
+                db, schema, stage, compute_pool
+            )
+
+            # upload the spec file to stage
+            result = self.session.file.put(snowpark_spec_file.name, "@%s" % stage)
+
+            service_name = name.replace("-", "_")
+            external_access = (
+                "EXTERNAL_ACCESS_INTEGRATIONS=(%s) " % ",".join(external_integration)
+                if external_integration
+                else ""
+            )
+
+            # cannot pass 'is_job' parameter using the API, thus we need to use SQL directly..
+            query = """
+            EXECUTE JOB SERVICE IN COMPUTE POOL {compute_pool}
+            NAME = {db}.{schema}.{service_name}
+            {external_access}
+            FROM @{stage} SPECIFICATION_FILE={specification_file}
+            """.format(
+                compute_pool=compute_pool,
+                db=db,
+                schema=schema,
+                service_name=service_name,
+                external_access=external_access,
+                stage=stage,
+                specification_file=result[0].target,
+            )
+
+            async_job = self.session.sql(query).collect(block=False)
+            return async_job.query_id, service_name
+
+    def terminate_job(self, service):
+        service.delete()

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py

@@ -0,0 +1,273 @@
+import os
+import sys
+import platform
+
+from metaflow import R, current
+from metaflow.metadata_provider import MetaDatum
+from metaflow.metadata_provider.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\n"
+                "Install required Snowflake packages using the @pypi decorator:\n"
+                "@pypi(packages={\n"
+                "    'snowflake': '1.8.0',\n"
+                "    'snowflake-connector-python': '3.18.0',\n"
+                "    'snowflake-snowpark-python': '1.40.0'\n"
+                "})\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,
+        "integration": None,  # Outerbounds OAuth integration name
+    }
+
+    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)
+
+        # Set defaults from config (user can override via decorator or integration)
+        if not self.attributes["account"]:
+            self.attributes["account"] = SNOWPARK_ACCOUNT
+        if not self.attributes["user"]:
+            self.attributes["user"] = SNOWPARK_USER
+        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
+        # Only use password from config if not using integration (OAuth)
+        if not self.attributes["integration"] and not self.attributes["password"]:
+            self.attributes["password"] = SNOWPARK_PASSWORD
+
+        # 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 packages first:\n"
+                "  snowflake==1.8.0\n"
+                "  snowflake-connector-python==3.18.0\n"
+                "  snowflake-snowpark-python==1.40.0\n\n"
+                "You can install them by executing:\n"
+                "%s -m pip install snowflake==1.8.0 snowflake-connector-python==3.18.0 snowflake-snowpark-python==1.40.0\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]