datatailr 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

datatailr/sbin/datatailr_run_app.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+
+# *************************************************************************
+#
+#  Copyright (c) 2025 - Datatailr Inc.
+#  All Rights Reserved.
+#
+#  This file is part of Datatailr and subject to the terms and conditions
+#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
+#  of this file, in parts or full, via any medium is strictly prohibited.
+# *************************************************************************
+
+import os
+
+from datatailr.logging import DatatailrLogger
+
+logger = DatatailrLogger(os.path.abspath(__file__)).get_logger()
+
+
+def run():
+    logger.info("Starting Datatailr app...")
+    entrypoint = os.environ.get("DATATAILR_ENTRYPOINT")
+
+    if entrypoint is None:
+        raise ValueError("Environment variable 'DATATAILR_ENTRYPOINT' is not set.")
+
+    os.system(entrypoint)
+    logger.info(f"Running entrypoint: {entrypoint}")

datatailr/sbin/{run_job.py → datatailr_run_batch.py}

@@ -12,19 +12,19 @@
 
 import importlib
 import os
-import pickle
 
-from datatailr import dt__Blob
 from datatailr.logging import DatatailrLogger
 
 logger = DatatailrLogger(os.path.abspath(__file__)).get_logger()
 
 
-def main():
+def run():
+    logger.info("Running Datatailr batch job")
     entry_point = os.environ.get("DATATAILR_BATCH_ENTRYPOINT")
     batch_run_id = os.environ.get("DATATAILR_BATCH_RUN_ID")
     batch_id = os.environ.get("DATATAILR_BATCH_ID")
     job_id = os.environ.get("DATATAILR_JOB_ID")
+    logger.info(f"Batch run ID: {batch_run_id}, Batch ID: {batch_id}, Job ID: {job_id}")
 
     if entry_point is None:
         raise ValueError(
@@ -44,20 +44,5 @@ def main():
         raise ValueError(
             f"The function '{func_name}' in module '{module_name}' is not callable."
         )
-    result = function()
-    result_path = f"batch-results-{batch_run_id}-{job_id}.pkl"
-    with open(result_path, "wb") as f:
-        pickle.dump(result, f)
-    blob = dt__Blob()
-    blob.cp(result_path, "blob://")
-    logger.info(f"{result_path} copied to blob storage.")
-
-
-if __name__ == "__main__":
-    try:
-        logger.debug("Starting job execution...")
-        main()
-        logger.debug("Job executed successfully.")
-    except Exception as e:
-        logger.error(f"Error during job execution: {e}")
-        raise
+    function()
+    logger.info("Datatailr batch job completed successfully.")

datatailr/scheduler/__init__.py

@@ -8,31 +8,47 @@
 #  of this file, in parts or full, via any medium is strictly prohibited.
 # *************************************************************************
 
-from datatailr.errors import BatchJobError, DatatailrError
+r"""
+Datatailr Scheduler Module
+==========================
+
+The `datatailr.scheduler` module provides a framework for scheduling and managing batch jobs.
+
+The main job types are:
+_______________________
+
+- **Batch**: Represents a batch job that can be scheduled and executed.
+  The job can include multiple tasks which can be run in parallel or sequentially.
+- **Service**: Represents a service job that runs continuously.
+- **App**: Represents a web app or a dashboard, which can be built using one of the supported frameworks,
+  such as `Streamlit <https://streamlit.io/>`_, `Dash <https://dash.plotly.com/>`_, or `Panel <https://panel.holoviz.org/>`_.
+- **Excel**: Represents an Excel add-in.
+"""
+
+from datatailr.errors import BatchJobError
 from datatailr.scheduler.base import (
-    ACL,
     EntryPoint,
     Environment,
     Job,
     JobType,
     Resources,
-    User,
+    set_allow_unsafe_scheduling,
 )
 from datatailr.scheduler.batch import Batch, BatchJob, DuplicateJobNameError
-from datatailr.scheduler.batch_decorator import batch_decorator as batch
+from datatailr.scheduler.batch_decorator import batch_decorator as batch_job
+from datatailr.scheduler.schedule import Schedule
 
 __all__ = [
     "Job",
     "JobType",
     "Environment",
-    "User",
     "Resources",
-    "ACL",
     "EntryPoint",
     "Batch",
     "BatchJob",
-    "batch",
-    "DatatailrError",
+    "batch_job",
     "BatchJobError",
     "DuplicateJobNameError",
+    "set_allow_unsafe_scheduling",
+    "Schedule",
 ]

datatailr/scheduler/arguments_cache.py

@@ -21,17 +21,26 @@ and the inner dictionaries contain the arguments.
 This module is for internal use of the datatailr package.
 """
 
-from collections import defaultdict
+from datatailr.dt_json import json, decode_json
+import os
 import pickle
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from datatailr import is_dt_installed, Blob
-from datatailr.scheduler import BatchJob
+from datatailr.errors import DatatailrError
 
 
 __BLOB_STORAGE__ = Blob()
 
 
+class CacheNotFoundError(DatatailrError):
+    """Custom error for cache operations."""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+        self.message = message
+
+
 class ArgumentsCache:
     def __init__(self, use_persistent_cache: bool = is_dt_installed()):
         """
@@ -40,11 +49,12 @@ class ArgumentsCache:
         :param use_persistent_cache: If True, use the persistent cache backend. Otherwise, use in-memory cache.
         """
         self.use_persistent_cache = use_persistent_cache
-        self.in_memory_cache: Dict[str, Dict[str, Dict[str, Any]]] = defaultdict(
-            lambda: defaultdict(dict)
-        )
+        if not self.use_persistent_cache:
+            # Create a temp folder, for local caching
+            os.makedirs("/tmp/datatailr/batch/arguments", exist_ok=True)
+            os.makedirs("/tmp/datatailr/batch/results", exist_ok=True)
 
-    def add_arguments(self, batch_run_id: str, job: str, arguments: Dict[str, Any]):
+    def add_arguments(self, batch_id: str, arguments: Dict[str, Any]):
         """
         Add arguments to the cache for a specific job and batch run.
 
@@ -52,13 +62,16 @@ class ArgumentsCache:
         :param job_name: Name of the job.
         :param arguments: Dictionary of arguments to store.
         """
-        if self.use_persistent_cache and isinstance(job, str):
-            path = f"{batch_run_id}/{job}/args"
+        path = f"/tmp/datatailr/batch/arguments/{batch_id}.pkl"
+        if self.use_persistent_cache:
             self._add_to_persistent_cache(path, arguments)
         else:
-            self.in_memory_cache[batch_run_id][job]["args"] = arguments
+            with open(path, "wb") as f:
+                pickle.dump(arguments, f)
 
-    def get_arguments(self, batch_run_id: str, job: str) -> Dict[str, Any]:
+    def get_arguments(
+        self, batch_id: str, job: str, batch_run_id: Optional[str]
+    ) -> Dict[str, Any]:
         """
         Retrieve arguments from the cache for a specific job and batch run.
 
@@ -66,27 +79,37 @@ class ArgumentsCache:
         :param job_name: Name of the job.
         :return: Dictionary of arguments.
         """
+        path = f"/tmp/datatailr/batch/arguments/{batch_id}.pkl"
         if self.use_persistent_cache and isinstance(job, str):
-            path = f"{batch_run_id}/{job}/args"
-            arg_keys = self._get_from_persistent_cache(path)
-            if not isinstance(arg_keys, dict):
-                raise TypeError(
-                    f"Expected a dictionary for arguments, got {type(arg_keys)}"
-                )
+            try:
+                arg_keys = self._get_from_persistent_cache(path)
+            except RuntimeError:
+                return {}
         else:
-            arg_keys = (
-                self.in_memory_cache.get(batch_run_id, {})
-                .get(job, {})
-                .get("args", {})
-                .items()
-            )
-        arguments = {}
-        for key, value in arg_keys:
-            if isinstance(value, BatchJob):
-                arguments[key] = value.name
-            else:
-                arguments[key] = value
-        return arguments
+            if not os.path.exists(path):
+                raise CacheNotFoundError(
+                    f"Cache file not found: {path}. Ensure that the arguments have been cached."
+                )
+            with open(path, "rb") as f:
+                try:
+                    arg_keys = pickle.load(f)
+                except EOFError:
+                    return {}
+                if not isinstance(arg_keys, dict):
+                    raise TypeError(
+                        f"Expected a dictionary for arguments, got {type(arg_keys)}"
+                    )
+        if batch_run_id is None:
+            return arg_keys[job]
+        arguments_mapping = decode_json(
+            os.getenv("DATATAILR_JOB_ARGUMENT_MAPPING", "{}")
+        )
+        arguments_mapping = {value: key for key, value in arguments_mapping.items()}
+        args = {
+            arguments_mapping.get(name, name): self.get_result(batch_run_id, value)
+            for name, value in arg_keys[job].items()
+        }
+        return args
 
     def add_result(self, batch_run_id: str, job: str, result: Any):
         """
@@ -96,13 +119,14 @@ class ArgumentsCache:
         :param job: Name of the job.
         :param result: Result of the batch job.
         """
+        path = f"/tmp/datatailr/batch/results/{batch_run_id}_{job}.pkl"
         if self.use_persistent_cache and isinstance(job, str):
-            path = f"{batch_run_id}/{job}/result"
             self._add_to_persistent_cache(path, result)
         else:
-            self.in_memory_cache[batch_run_id][job]["result"] = result
+            with open(path, "wb") as f:
+                pickle.dump(result, f)
 
-    def get_result(self, batch_run_id: str, job: str) -> Any:
+    def get_result(self, batch_run_id: str, job: Any) -> Any:
         """
         Retrieve the result of a batch job from the cache.
 
@@ -110,10 +134,17 @@ class ArgumentsCache:
         :param job: Name of the job.
         :return: Result of the batch job.
         """
+        path = f"/tmp/datatailr/batch/results/{batch_run_id}_{job}.pkl"
         if self.use_persistent_cache and isinstance(job, str):
-            path = f"{batch_run_id}/{job}/result"
             return self._get_from_persistent_cache(path)
-        return self.in_memory_cache[batch_run_id][job].get("result")
+        else:
+            if not os.path.exists(path):
+                return job
+            with open(path, "rb") as f:
+                try:
+                    return pickle.load(f)
+                except EOFError:
+                    return None
 
     def _add_to_persistent_cache(self, path: str, blob: Any):
         """
@@ -124,9 +155,8 @@ class ArgumentsCache:
         :raises TypeError: If the blob cannot be pickled.
 
         """
-        __BLOB_STORAGE__.put_blob(
-            path, pickle.dumps(blob, protocol=pickle.HIGHEST_PROTOCOL)
-        )
+        path = path.replace("/tmp/", "")
+        __BLOB_STORAGE__.put_blob(path, json.dumps(blob))
 
     def _get_from_persistent_cache(self, path: str) -> Any:
         """
@@ -134,8 +164,6 @@ class ArgumentsCache:
 
         :param path: Path in the Blob storage where the blob is stored.
         """
-        try:
-            data = __BLOB_STORAGE__.get_blob(path)
-            return pickle.loads(data)
-        except (TypeError, EOFError):
-            return {}
+        path = path.replace("/tmp/", "")
+        data = __BLOB_STORAGE__.get_blob(path)
+        return json.loads(data)

datatailr/scheduler/base.py

@@ -8,7 +8,11 @@
 #  of this file, in parts or full, via any medium is strictly prohibited.
 # *************************************************************************
 
+from __future__ import annotations
+
+from datetime import datetime
 import importlib
+import inspect
 import json
 import os
 import subprocess
@@ -18,12 +22,27 @@ from dataclasses import dataclass
 from enum import Enum
 from typing import Callable, Optional, Tuple, Union
 
-from datatailr import ACL, Environment, User, dt__Job, is_dt_installed
+from datatailr import ACL, Environment, User, is_dt_installed
+from datatailr.wrapper import dt__Job
+from datatailr.scheduler.constants import DEFAULT_TASK_MEMORY, DEFAULT_TASK_CPU
 from datatailr.build.image import Image
 from datatailr.errors import BatchJobError
 from datatailr.logging import DatatailrLogger
+from datatailr.utils import run_shell_command
 
 logger = DatatailrLogger(os.path.abspath(__file__)).get_logger()
+__client__ = dt__Job()
+
+
+def set_allow_unsafe_scheduling(allow: bool):
+    """
+    Set whether to allow unsafe scheduling of jobs.
+    This is a global setting that affects how jobs are scheduled.
+    """
+    if allow:
+        os.environ["DATATAILR_ALLOW_UNSAFE_SCHEDULING"] = "true"
+    else:
+        os.environ.pop("DATATAILR_ALLOW_UNSAFE_SCHEDULING", None)
 
 
 class RepoValidationError(BatchJobError):
@@ -40,6 +59,7 @@ class JobType(Enum):
     BATCH = "batch"
     SERVICE = "service"
     APP = "app"
+    EXCEL = "excel"
     UNKNOWN = "unknown"
 
     def __str__(self):
@@ -55,8 +75,14 @@ class Resources:
     Represents the resources required for a job.
     """
 
-    memory: str = "100m"
-    cpu: int = 1
+    memory: str = DEFAULT_TASK_MEMORY
+    cpu: float = DEFAULT_TASK_CPU
+
+
+# TODO: create a dt_run script that will:
+# 1. create user and group if not exists
+# 2. set the correct path
+# 3. run the job based on its type
 
 
 class EntryPoint:
@@ -68,26 +94,30 @@ class EntryPoint:
     def __init__(
         self,
         type: JobType,
-        func: Optional[Callable] = None,
-        module_name: Optional[str] = None,
-        function_name: Optional[str] = None,
+        func: Callable,
     ):
-        if func is None and (module_name is None or function_name is None):
-            raise ValueError(
-                "Either a function or module and function names must be provided."
-            )
         self.func = func
-        self.module_name = func.__module__ if func else module_name
-        self.function_name = func.__name__ if func else function_name
+        self.module_name = func.__module__
+        self.function_name = func.__name__
         self.type = type
 
+        # Find the absolute path to the repository and then the relative path to the module.
+        # This will be used in the creation of the code 'bundle' when building the image.
+        path_to_repo = run_shell_command("git rev-parse --show-toplevel")[0]
+        path_to_code = inspect.getfile(func)
+        package_root = path_to_code
+        module_parts = self.module_name.split(".")
+        for _ in module_parts:
+            package_root = os.path.dirname(package_root)
+        path_to_module = os.path.relpath(package_root, path_to_repo)
+        self.path_to_repo = path_to_repo
+        self.path_to_module = path_to_module
+
     def __call__(self, *args, **kwargs):
+        os.environ.update(kwargs.pop("env", {}))
         if self.type == JobType.BATCH:
-            if self.module_name and self.function_name:
-                module = importlib.import_module(self.module_name)
-                func = getattr(module, self.function_name)
-            elif self.func is not None:
-                func = self.func
+            module = importlib.import_module(self.module_name)
+            func = getattr(module, self.function_name)
             return func(*args, **kwargs)
 
         elif self.type == JobType.SERVICE:
@@ -106,13 +136,28 @@ class EntryPoint:
 class Job:
     def __init__(
         self,
-        environment: Optional[Environment],
         name: str,
-        image: Image,
-        run_as: Optional[Union[str, User]],
+        environment: Optional[Environment] = Environment.DEV,
+        image: Optional[Image] = None,
+        run_as: Optional[Union[str, User]] = User.signed_user(),
         resources: Resources = Resources(memory="100m", cpu=1),
         acl: Optional[ACL] = None,
+        python_requirements: str = "",
+        build_script_pre: str = "",
+        build_script_post: str = "",
+        type: JobType = JobType.UNKNOWN,
+        entrypoint: Optional[EntryPoint] = None,
+        update_existing: bool = False,
     ):
+        if environment is None:
+            environment = Environment.DEV
+
+        if update_existing:
+            existing_job = self.__get_existing__(name, environment)
+            if existing_job:
+                self.from_dict(existing_job)
+                return
+
         if run_as is None:
             run_as = User.signed_user()
         if environment is None:
@@ -126,11 +171,16 @@ class Job:
         self.name = name
         self.run_as = run_as
         self.resources = resources
+        if image is None:
+            image = Image(
+                acl=self.acl,
+                python_requirements=python_requirements,
+                build_script_pre=build_script_pre,
+                build_script_post=build_script_post,
+            )
         self.image = image
-
-        # Placeholders, to be set in derived classes
-        self.type: JobType = JobType.UNKNOWN
-        self.entrypoint = None
+        self.type = type
+        self.entrypoint = entrypoint
         self.__id = str(uuid.uuid4())
 
     @property
@@ -140,6 +190,25 @@ class Job:
         """
         return self.__id
 
+    @classmethod
+    def __get_existing__(
+        cls, job_name: str, environment: Environment
+    ) -> Optional[dict]:
+        """
+        Retrieve an existing job instance from the DataTailr platform.
+        Based on the job name and environment.
+        """
+        job_list = __client__.ls(filter=f"name={job_name},environment={environment}")
+        if not isinstance(job_list, list):
+            return None
+        if len(job_list) == 0:
+            return None
+        if len(job_list) > 1:
+            raise BatchJobError(
+                f"Multiple jobs found with name '{job_name}' in environment '{environment}'."
+            )
+        return job_list[0]
+
     def __repr__(self):
         return (
             f"Job(name={self.name}, environment={self.environment}, "
@@ -169,31 +238,64 @@ class Job:
             job_dict["cpu"] = self.resources.cpu
         return job_dict
 
+    def from_dict(self, job_dict: dict):
+        self.name = job_dict["name"]
+        self.image = job_dict["image"]
+
+        environment = job_dict.get("environment", "dev")
+        environment = Environment(environment.lower())
+        self.environment = environment
+
+        user = job_dict["run_as"]["name"]
+        user = User(user.lower())
+        self.run_as = user
+
+        self.resources = Resources(memory=job_dict["memory"], cpu=job_dict["num_cpus"])
+        acl = job_dict.get("acl", None)
+        if acl is None:
+            acl = ACL(user=self.run_as)
+        else:
+            acl = ACL.from_dict(acl)
+        self.acl = acl
+        self.python_requirements = (job_dict.get("python_requirements", ""),)
+        self.build_script_pre = (job_dict.get("build_script_pre", ""),)
+        self.build_script_post = (job_dict.get("build_script_post", ""),)
+        self.type = JobType(job_dict.get("type", "unknown").lower())
+        self.state = job_dict["state"]
+        self.create_time = datetime.fromtimestamp(job_dict["create_time"] * 1e-6)
+        self.version = job_dict["version"]
+        self.__id = job_dict["id"]
+
     def to_json(self):
         """
         Convert the Job instance to a JSON string representation.
         """
         return json.dumps(self.to_dict())
 
-    def verify_repo_is_ready(self) -> Tuple[bool, str]:
-        is_committed = (
-            subprocess.run(
-                ("git diff --exit-code"), shell=True, capture_output=True
-            ).returncode
-            == 0
-        )
+    def verify_repo_is_ready(self) -> Tuple[str, str]:
+        """
+        Verify if the repository is ready for job execution.
+        The check consists of two parts:
+        1. Check if there are uncommitted changes in the repository.
+        2. Check if the local commit matches the remote HEAD (the repo is synced with the remote).
+        Returns a tuple of (branch: str, commit_hash: str).
+        """
+        local_commit = run_shell_command("git rev-parse HEAD")[0]
+        branch_name = run_shell_command("git rev-parse --abbrev-ref HEAD")[0]
+
+        if os.getenv("DATATAILR_ALLOW_UNSAFE_SCHEDULING", "false").lower() == "true":
+            return branch_name, local_commit
+        return_code = run_shell_command("git diff --exit-code")[1]
+        is_committed = return_code == 0
+
         if not is_committed:
-            return (
-                False,
-                "Uncommitted changes detected. Please commit your changes before running the job.",
+            raise RepoValidationError(
+                "Please commit your changes before running the job."
             )
 
-        local_commit = subprocess.run(
-            ("git rev-parse HEAD"), shell=True, capture_output=True, text=True
-        ).stdout.strip()
         remote_commit = (
             subprocess.run(
-                ("git ls-remote origin HEAD"),
+                ("remote_commit = $(git ls-remote origin HEAD)"),
                 shell=True,
                 capture_output=True,
                 text=True,
@@ -203,43 +305,67 @@ class Job:
         )
 
         if local_commit != remote_commit:
-            return (
-                False,
-                "Local commit does not match remote HEAD. Please pull the latest changes before running the job.",
+            raise RepoValidationError(
+                "Please sync your local repository with the remote before running the job."
             )
 
-        branch = subprocess.run(
-            ("git rev-parse --abbrev-ref HEAD"),
-            shell=True,
-            capture_output=True,
-            text=True,
-        ).stdout.strip()
-        return True, ""
+        return branch_name, local_commit
 
-    def run(self) -> Tuple[bool, str]:
-        """
-        Run the job. This method should be implemented to execute the job logic.
-        It verifies the repository state and prepares the job for execution.
-        Returns a tuple of (success: bool, message: str).
-        If the repository is not ready, it returns False with an error message.
-        If the job runs successfully, it returns True with an empty message.
-        """
-        if is_dt_installed():
-            check_result = self.verify_repo_is_ready()
-            if not check_result[0]:
-                raise RepoValidationError(check_result[1])
-            logger.info(
-                f"Running job '{self.name}' in environment '{self.environment}' as '{self.run_as}'"
-            )
+    def __prepare__(self) -> str:
+        branch_name, local_commit = self.verify_repo_is_ready()
+        self.image.update(
+            branch_name=branch_name,
+            commit_hash=local_commit,
+        )
+        logger.info(
+            f"Running job '{self.name}' in environment '{self.environment}' as '{self.run_as}'"
+        )
 
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".json") as temp_file:
-                temp_file.write(self.to_json().encode())
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".json") as temp_file:
+            temp_file.write(self.to_json().encode())
+        return temp_file.name
 
-            dt__Job().run(f"file://{temp_file.name}")
-            os.remove(temp_file.name)
+    def get_schedule_args(self) -> dict:
+        """
+        Returns additional arguments for scheduling the job.
+        Override or extend this method as needed.
+        """
+        return {}
 
-            return True, ""
-        else:
+    def __run_command__(self, command: str) -> Tuple[bool, str]:
+        """
+        Run a command in the context of the job.
+        This is used to execute the job's entry point.
+        """
+        if not is_dt_installed():
             raise NotImplementedError(
                 "DataTailr is not installed. Please install DataTailr to run this job."
             )
+        try:
+            temp_file_name = self.__prepare__()
+
+            if command == "run":
+                __client__.run(f"file://{temp_file_name}", **self.get_schedule_args())
+            elif command == "save":
+                __client__.save(f"file://{temp_file_name}", **self.get_schedule_args())
+            else:
+                raise ValueError(f"Unknown command: {command}")
+            os.remove(temp_file_name)
+        except Exception as e:
+            logger.error(f"Error running command '{command}': {e}")
+            return False, str(e)
+        return True, f"Job '{self.name}' {command}d successfully."
+
+    def save(self) -> Tuple[bool, str]:
+        """
+        Save the job to the DataTailr platform.
+        If the job already exists, it will be updated.
+        """
+        return self.__run_command__("save")
+
+    def run(self) -> Tuple[bool, str]:
+        """
+        Run the job. This method should be implemented to execute the job logic.
+        It verifies the repository state and prepares the job for execution.
+        """
+        return self.__run_command__("run")