ob-metaflow-extensions 1.1.79__py2.py3-none-any.whl → 1.1.81__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py

@@ -0,0 +1,299 @@
+import os
+import shlex
+import atexit
+import json
+import math
+import time
+
+from metaflow import util
+
+from metaflow.metaflow_config import (
+    SERVICE_INTERNAL_URL,
+    SERVICE_HEADERS,
+    DEFAULT_METADATA,
+    DATASTORE_SYSROOT_S3,
+    DATATOOLS_S3ROOT,
+    KUBERNETES_SANDBOX_INIT_SCRIPT,
+    OTEL_ENDPOINT,
+    DEFAULT_SECRETS_BACKEND_TYPE,
+    AWS_SECRETS_MANAGER_DEFAULT_REGION,
+    S3_SERVER_SIDE_ENCRYPTION,
+    S3_ENDPOINT_URL,
+)
+from metaflow.metaflow_config_funcs import config_values
+
+from metaflow.mflog import (
+    export_mflog_env_vars,
+    bash_capture_logs,
+    BASH_SAVE_LOGS,
+    get_log_tailer,
+    tail_logs,
+)
+
+from .snowpark_client import SnowparkClient
+from .snowpark_exceptions import SnowparkException
+from .snowpark_job import SnowparkJob
+
+# Redirect structured logs to $PWD/.logs/
+LOGS_DIR = "$PWD/.logs"
+STDOUT_FILE = "mflog_stdout"
+STDERR_FILE = "mflog_stderr"
+STDOUT_PATH = os.path.join(LOGS_DIR, STDOUT_FILE)
+STDERR_PATH = os.path.join(LOGS_DIR, STDERR_FILE)
+
+
+class Snowpark(object):
+    def __init__(
+        self,
+        datastore,
+        metadata,
+        environment,
+        client_credentials,
+    ):
+        self.datastore = datastore
+        self.metadata = metadata
+        self.environment = environment
+        self.snowpark_client = SnowparkClient(**client_credentials)
+        atexit.register(lambda: self.job.kill() if hasattr(self, "job") else None)
+
+    def _job_name(self, user, flow_name, run_id, step_name, task_id, retry_count):
+        return "{user}-{flow_name}-{run_id}-{step_name}-{task_id}-{retry_count}".format(
+            user=user,
+            flow_name=flow_name,
+            run_id=str(run_id) if run_id is not None else "",
+            step_name=step_name,
+            task_id=str(task_id) if task_id is not None else "",
+            retry_count=str(retry_count) if retry_count is not None else "",
+        )
+
+    def _command(self, environment, code_package_url, step_name, step_cmds, task_spec):
+        mflog_expr = export_mflog_env_vars(
+            datastore_type=self.datastore.TYPE,
+            stdout_path=STDOUT_PATH,
+            stderr_path=STDERR_PATH,
+            **task_spec
+        )
+        init_cmds = environment.get_package_commands(
+            code_package_url, self.datastore.TYPE
+        )
+        init_expr = " && ".join(init_cmds)
+        step_expr = bash_capture_logs(
+            " && ".join(
+                environment.bootstrap_commands(step_name, self.datastore.TYPE)
+                + step_cmds
+            )
+        )
+
+        # construct an entry point that
+        # 1) initializes the mflog environment (mflog_expr)
+        # 2) bootstraps a metaflow environment (init_expr)
+        # 3) executes a task (step_expr)
+
+        # the `true` command is to make sure that the generated command
+        # plays well with docker containers which have entrypoint set as
+        # eval $@
+        cmd_str = "true && mkdir -p %s && %s && %s && %s; " % (
+            LOGS_DIR,
+            mflog_expr,
+            init_expr,
+            step_expr,
+        )
+        # after the task has finished, we save its exit code (fail/success)
+        # and persist the final logs. The whole entrypoint should exit
+        # with the exit code (c) of the task.
+        #
+        # Note that if step_expr OOMs, this tail expression is never executed.
+        # We lose the last logs in this scenario.
+        cmd_str += "c=$?; %s; exit $c" % BASH_SAVE_LOGS
+        # For supporting sandboxes, ensure that a custom script is executed before
+        # anything else is executed. The script is passed in as an env var.
+        cmd_str = (
+            '${METAFLOW_INIT_SCRIPT:+eval \\"${METAFLOW_INIT_SCRIPT}\\"} && %s'
+            % cmd_str
+        )
+        return shlex.split('bash -c "%s"' % cmd_str)
+
+    def create_job(
+        self,
+        step_name,
+        step_cli,
+        task_spec,
+        code_package_sha,
+        code_package_url,
+        code_package_ds,
+        image=None,
+        stage=None,
+        compute_pool=None,
+        volume_mounts=None,
+        external_integration=None,
+        cpu=None,
+        gpu=None,
+        memory=None,
+        run_time_limit=None,
+        env=None,
+        attrs=None,
+    ) -> SnowparkJob:
+        if env is None:
+            env = {}
+        if attrs is None:
+            attrs = {}
+
+        job_name = self._job_name(
+            attrs.get("metaflow.user"),
+            attrs.get("metaflow.flow_name"),
+            attrs.get("metaflow.run_id"),
+            attrs.get("metaflow.step_name"),
+            attrs.get("metaflow.task_id"),
+            attrs.get("metaflow.retry_count"),
+        )
+
+        snowpark_job = (
+            SnowparkJob(
+                client=self.snowpark_client,
+                name=job_name,
+                command=self._command(
+                    self.environment, code_package_url, step_name, [step_cli], task_spec
+                ),
+                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=code_package_ds,
+                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,
+            )
+            .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
+            .environment_variable("METAFLOW_CODE_URL", code_package_url)
+            .environment_variable("METAFLOW_CODE_DS", code_package_ds)
+            .environment_variable("METAFLOW_USER", attrs["metaflow.user"])
+            .environment_variable("METAFLOW_SERVICE_URL", SERVICE_INTERNAL_URL)
+            .environment_variable(
+                "METAFLOW_SERVICE_HEADERS", json.dumps(SERVICE_HEADERS)
+            )
+            .environment_variable("METAFLOW_DATASTORE_SYSROOT_S3", DATASTORE_SYSROOT_S3)
+            .environment_variable("METAFLOW_DATATOOLS_S3ROOT", DATATOOLS_S3ROOT)
+            .environment_variable("METAFLOW_DEFAULT_DATASTORE", self.datastore.TYPE)
+            .environment_variable("METAFLOW_DEFAULT_METADATA", DEFAULT_METADATA)
+            .environment_variable("METAFLOW_SNOWPARK_WORKLOAD", 1)
+            .environment_variable("METAFLOW_RUNTIME_ENVIRONMENT", "snowpark")
+            .environment_variable(
+                "METAFLOW_INIT_SCRIPT", KUBERNETES_SANDBOX_INIT_SCRIPT
+            )
+            .environment_variable("METAFLOW_OTEL_ENDPOINT", OTEL_ENDPOINT)
+            .environment_variable(
+                "SNOWFLAKE_WAREHOUSE",
+                self.snowpark_client.connection_parameters.get("warehouse"),
+            )
+        )
+
+        for k, v in config_values():
+            if k.startswith("METAFLOW_CONDA_") or k.startswith("METAFLOW_DEBUG_"):
+                snowpark_job.environment_variable(k, v)
+
+        if DEFAULT_SECRETS_BACKEND_TYPE is not None:
+            snowpark_job.environment_variable(
+                "METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE", DEFAULT_SECRETS_BACKEND_TYPE
+            )
+        if AWS_SECRETS_MANAGER_DEFAULT_REGION is not None:
+            snowpark_job.environment_variable(
+                "METAFLOW_AWS_SECRETS_MANAGER_DEFAULT_REGION",
+                AWS_SECRETS_MANAGER_DEFAULT_REGION,
+            )
+        if S3_SERVER_SIDE_ENCRYPTION is not None:
+            snowpark_job.environment_variable(
+                "METAFLOW_S3_SERVER_SIDE_ENCRYPTION", S3_SERVER_SIDE_ENCRYPTION
+            )
+        if S3_ENDPOINT_URL is not None:
+            snowpark_job.environment_variable(
+                "METAFLOW_S3_ENDPOINT_URL", S3_ENDPOINT_URL
+            )
+
+        for name, value in env.items():
+            snowpark_job.environment_variable(name, value)
+
+        return snowpark_job
+
+    def launch_job(self, **kwargs):
+        self.job = self.create_job(**kwargs).create().execute()
+
+    def wait(self, stdout_location, stderr_location, echo=None):
+        def update_delay(secs_since_start):
+            # this sigmoid function reaches
+            # - 0.1 after 11 minutes
+            # - 0.5 after 15 minutes
+            # - 1.0 after 23 minutes
+            # in other words, the user will see very frequent updates
+            # during the first 10 minutes
+            sigmoid = 1.0 / (1.0 + math.exp(-0.01 * secs_since_start + 9.0))
+            return 0.5 + sigmoid * 30.0
+
+        def wait_for_launch(job):
+            status = job.status
+
+            echo(
+                "Task is starting (%s)..." % status,
+                "stderr",
+                job_id=job.id,
+            )
+            t = time.time()
+            start_time = time.time()
+            while job.is_waiting:
+                new_status = job.status
+                if status != new_status or (time.time() - t) > 30:
+                    status = new_status
+                    echo(
+                        "Task is starting (%s)..." % status,
+                        "stderr",
+                        job_id=job.id,
+                    )
+                    t = time.time()
+                time.sleep(update_delay(time.time() - start_time))
+
+        _make_prefix = lambda: b"[%s] " % util.to_bytes(self.job.id)
+
+        stdout_tail = get_log_tailer(stdout_location, self.datastore.TYPE)
+        stderr_tail = get_log_tailer(stderr_location, self.datastore.TYPE)
+
+        # 1) Loop until the job has started
+        wait_for_launch(self.job)
+
+        # 2) Tail logs until the job has finished
+        tail_logs(
+            prefix=_make_prefix(),
+            stdout_tail=stdout_tail,
+            stderr_tail=stderr_tail,
+            echo=echo,
+            has_log_updates=lambda: self.job.is_running,
+        )
+
+        if self.job.has_failed:
+            msg = next(
+                msg
+                for msg in [
+                    self.job.message,
+                    "Task crashed.",
+                ]
+                if msg is not None
+            )
+            raise SnowparkException(
+                "%s " "This could be a transient error. " "Use @retry to retry." % msg
+            )
+        else:
+            if self.job.is_running:
+                # Kill the job if it is still running by throwing an exception.
+                raise SnowparkException("Task failed!")
+            echo(
+                "Task finished with message '%s'." % self.job.message,
+                "stderr",
+                job_id=self.job.id,
+            )

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py

@@ -0,0 +1,271 @@
+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.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(
+    "--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,
+    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,
+            },
+        )
+        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,123 @@
+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,
+        user: str,
+        password: str,
+        role: str,
+        database: str,
+        warehouse: str,
+        schema: str,
+        autocommit: bool = True,
+    ):
+        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 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
+            )
+
+        self.connection_parameters = {
+            "account": account,
+            "user": user,
+            "password": password,
+            "role": role,
+            "warehouse": warehouse,
+            "database": database,
+            "schema": schema,
+            "autocommit": autocommit,
+        }
+
+        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()