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_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,241 @@
+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 dashes..
+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):
+        status_list = self.status_obj()
+        if not status_list:
+            return "UNKNOWN"
+        return status_list[0].get("status", "UNKNOWN")
+
+    @property
+    def message(self):
+        status_list = self.status_obj()
+        if not status_list:
+            return None
+        return status_list[0].get("message")
+
+    @property
+    def is_waiting(self):
+        return self.status in ["PENDING", "UNKNOWN"]
+
+    @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

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_service_spec.py

@@ -0,0 +1,259 @@
+import json
+from collections import defaultdict
+from typing import List, Dict, Optional
+
+
+class Container:
+    def __init__(self, name: str, image: str):
+        self.payload = defaultdict(lambda: defaultdict(dict))
+        self.payload["name"] = name
+        self.payload["image"] = image
+
+    def command(self, command: List[str]) -> "Container":
+        self.payload["command"] = command
+        return self
+
+    def args(self, args: List[str]) -> "Container":
+        self.payload["args"] = args
+        return self
+
+    def env(self, env: Dict[str, str]) -> "Container":
+        self.payload["env"] = env
+        return self
+
+    def readiness_probe(self, readiness_probe: "ReadinessProbe") -> "Container":
+        self.payload["readiness_probe"] = (
+            readiness_probe.to_dict() if readiness_probe else None
+        )
+        return self
+
+    def volume_mounts(self, volume_mounts: List["VolumeMount"]) -> "Container":
+        self.payload["volume_mounts"] = [vm.to_dict() for vm in volume_mounts]
+        return self
+
+    def resources(self, resources: "Resources") -> "Container":
+        self.payload["resources"] = resources.to_dict() if resources else None
+        return self
+
+    def secrets(self, secrets: List["Secrets"]) -> "Container":
+        self.payload["secrets"] = [secret.to_dict() for secret in secrets]
+        return self
+
+    def to_dict(self) -> Dict:
+        result = {"name": self.payload["name"], "image": self.payload["image"]}
+        if "command" in self.payload and self.payload["command"]:
+            result["command"] = self.payload["command"]
+        if "args" in self.payload and self.payload["args"]:
+            result["args"] = self.payload["args"]
+        if "env" in self.payload and self.payload["env"]:
+            result["env"] = self.payload["env"]
+        if "readiness_probe" in self.payload and self.payload["readiness_probe"]:
+            result["readiness_probe"] = self.payload["readiness_probe"]
+        if "volume_mounts" in self.payload and self.payload["volume_mounts"]:
+            result["volume_mounts"] = self.payload["volume_mounts"]
+        if "resources" in self.payload and self.payload["resources"]:
+            result["resources"] = self.payload["resources"]
+        if "secrets" in self.payload and self.payload["secrets"]:
+            result["secrets"] = self.payload["secrets"]
+        return result
+
+
+class ReadinessProbe:
+    def __init__(self, port: int, path: str):
+        self.payload = defaultdict(dict)
+        self.payload["port"] = port
+        self.payload["path"] = path
+
+    def to_dict(self) -> Dict:
+        return dict(self.payload)
+
+
+class VolumeMount:
+    def __init__(self, name: str, mount_path: str):
+        self.payload = defaultdict(dict)
+        self.payload["name"] = name
+        self.payload["mount_path"] = mount_path
+
+    def to_dict(self) -> Dict:
+        return dict(self.payload)
+
+
+class Resources:
+    def __init__(
+        self,
+        requests: Optional[Dict[str, str]] = None,
+        limits: Optional[Dict[str, str]] = None,
+    ):
+        self.payload = defaultdict(dict)
+        if requests:
+            self.payload["requests"] = requests
+        if limits:
+            self.payload["limits"] = limits
+
+    def to_dict(self) -> Dict:
+        result = {}
+        if "requests" in self.payload and self.payload["requests"]:
+            result["requests"] = self.payload["requests"]
+        if "limits" in self.payload and self.payload["limits"]:
+            result["limits"] = self.payload["limits"]
+        return result
+
+
+class Secrets:
+    def __init__(self, snowflake_secret: str):
+        self.payload = {"snowflake_secret": snowflake_secret}
+
+    def secret_key_ref(self, secret_key_ref: str) -> "Secrets":
+        self.payload["secret_key_ref"] = secret_key_ref
+        return self
+
+    def env_var_name(self, env_var_name: str) -> "Secrets":
+        self.payload["env_var_name"] = env_var_name
+        return self
+
+    def directory_path(self, directory_path: str) -> "Secrets":
+        self.payload["directory_path"] = directory_path
+        return self
+
+    def to_dict(self) -> Dict:
+        return {k: v for k, v in self.payload.items() if v}
+
+
+class Endpoint:
+    def __init__(self, name: str, port: int):
+        self.payload = defaultdict(dict)
+        self.payload["name"] = name
+        self.payload["port"] = port
+
+    def public(self, public: bool) -> "Endpoint":
+        self.payload["public"] = public
+        return self
+
+    def protocol(self, protocol: str) -> "Endpoint":
+        self.payload["protocol"] = protocol
+        return self
+
+    def to_dict(self) -> Dict:
+        return dict(self.payload)
+
+
+class Volume:
+    def __init__(self, name: str, source: str):
+        self.payload = defaultdict(dict)
+        self.payload["name"] = name
+        self.payload["source"] = source
+
+    def size(self, size: str) -> "Volume":
+        self.payload["size"] = size
+        return self
+
+    def block_config(self, block_config: Dict) -> "Volume":
+        self.payload["block_config"] = block_config
+        return self
+
+    def uid(self, uid: int) -> "Volume":
+        self.payload["uid"] = uid
+        return self
+
+    def gid(self, gid: int) -> "Volume":
+        self.payload["gid"] = gid
+        return self
+
+    def to_dict(self) -> Dict:
+        result = {"name": self.payload["name"], "source": self.payload["source"]}
+        if "size" in self.payload:
+            result["size"] = self.payload["size"]
+        if "block_config" in self.payload:
+            result["block_config"] = self.payload["block_config"]
+        if "uid" in self.payload:
+            result["uid"] = self.payload["uid"]
+        if "gid" in self.payload:
+            result["gid"] = self.payload["gid"]
+        return result
+
+
+class LogExporters:
+    def __init__(self):
+        self.payload = {}
+
+    def event_table_config(self, log_level: str) -> "LogExporters":
+        self.payload["eventTableConfig"] = {"logLevel": log_level}
+        return self
+
+    def to_dict(self) -> Dict:
+        return dict(self.payload)
+
+
+class ServiceRole:
+    def __init__(self, name: str):
+        self.payload = {"name": name}
+
+    def endpoints(self, endpoints: List[str]) -> "ServiceRole":
+        self.payload["endpoints"] = endpoints
+        return self
+
+    def to_dict(self) -> Dict:
+        return dict(self.payload)
+
+
+class SnowparkServiceSpec:
+    def __init__(self):
+        self.payload = defaultdict(lambda: defaultdict(list))
+
+    def containers(self, containers: List[Container]) -> "SnowparkServiceSpec":
+        self.payload["containers"] = [container.to_dict() for container in containers]
+        return self
+
+    def endpoints(self, endpoints: List[Endpoint]) -> "SnowparkServiceSpec":
+        self.payload["endpoints"] = [endpoint.to_dict() for endpoint in endpoints]
+        return self
+
+    def volumes(self, volumes: List[Volume]) -> "SnowparkServiceSpec":
+        self.payload["volumes"] = [volume.to_dict() for volume in volumes]
+        return self
+
+    def log_exporters(self, log_exporters: LogExporters) -> "SnowparkServiceSpec":
+        self.payload["logExporters"] = log_exporters.to_dict()
+        return self
+
+    def service_roles(self, service_roles: List[ServiceRole]) -> "SnowparkServiceSpec":
+        self.payload["serviceRoles"] = [role.to_dict() for role in service_roles]
+        return self
+
+    def to_dict(self) -> Dict:
+        return {"spec": dict(self.payload)}
+
+
+def generate_spec_file(spec: SnowparkServiceSpec, filename: str, format: str = "yaml"):
+    import yaml
+
+    spec_dict = spec.to_dict()
+    with open(filename, "w") as file:
+        if format == "json":
+            json.dump(spec_dict, file, indent=2)
+        elif format == "yaml":
+            yaml.dump(spec_dict, file, default_flow_style=False)
+
+
+if __name__ == "__main__":
+    # Example usage
+    container = (
+        Container(name="example-container", image="example-image")
+        .command(["python3", "app.py"])
+        .env({"ENV_VARIABLE": "value"})
+        .readiness_probe(ReadinessProbe(port=8080, path="/health"))
+        .resources(
+            Resources(requests={"memory": "2G", "cpu": "1"}, limits={"memory": "4G"})
+        )
+    )
+
+    endpoint = Endpoint(name="example-endpoint", port=8080).public(True)
+    volume = Volume(name="example-volume", source="local")
+    spec = (
+        SnowparkServiceSpec()
+        .containers([container])
+        .endpoints([endpoint])
+        .volumes([volume])
+    )
+
+    generate_spec_file(spec, "service_spec.yaml", format="yaml")

metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py

@@ -0,0 +1,50 @@
+import functools
+from metaflow.decorators import StepDecorator
+
+
+class TensorboardDecorator(StepDecorator):
+    name = "tensorboard"
+    defaults = {}
+
+    def task_decorate(
+        self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
+    ):
+        @functools.wraps(step_func)
+        def tb_wrapper():
+            import sys, os
+            from metaflow import metaflow_config, current
+
+            try:
+                from torch.utils.tensorboard import SummaryWriter
+            except:
+                print(
+                    "[@tensorboard] Torch and tensorboard not found - logging disabled!",
+                    file=sys.stderr,
+                )
+                step_func()
+            else:
+                tb_root = os.path.join(metaflow_config.DATATOOLS_S3ROOT, "tb")
+                pathspec = current.pathspec
+                try:
+                    log_dir = os.path.join(tb_root, current.project_flow_name, pathspec)
+                except:
+                    log_dir = os.path.join(tb_root, pathspec)
+                comps = log_dir[len(tb_root) + 1 :].split("/")
+                run_level = "/".join(comps[:-2])
+                flow_level = "/".join(comps[:-3])
+
+                print("[@tensorboard] -- INSPECTING RESULTS")
+                print(
+                    "[@tensorboard] -- Execute one of these commands on your workstation:"
+                )
+                print(f"[@tensorboard] Compare tasks of this run: obtb {run_level}")
+                print(f"[@tensorboard] Compare across runs: obtb {flow_level}")
+                writer = SummaryWriter(log_dir=log_dir)
+                setattr(flow, "obtb", writer)
+                try:
+                    step_func()
+                finally:
+                    writer.flush()
+                    delattr(flow, "obtb")
+
+        return tb_wrapper