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

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")

datatailr/scheduler/batch.py

@@ -14,9 +14,11 @@ import contextvars
 import json
 import os
 from functools import reduce
-from typing import Dict, List, Optional, Sequence, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
+import uuid
 
 from datatailr import Image
+from datatailr.dt_json import encode_json
 from datatailr.errors import BatchJobError
 from datatailr.logging import DatatailrLogger
 from datatailr.scheduler.base import (
@@ -29,9 +31,12 @@ from datatailr.scheduler.base import (
     User,
 )
 from datatailr.scheduler.constants import DEFAULT_TASK_CPU, DEFAULT_TASK_MEMORY
+from datatailr.scheduler.arguments_cache import ArgumentsCache
+from datatailr.scheduler.schedule import Schedule
 from datatailr.utils import is_dt_installed
 
 __DAG_CONTEXT__: contextvars.ContextVar = contextvars.ContextVar("dag_context")
+__ARGUMENTS_CACHE__ = ArgumentsCache()
 logger = DatatailrLogger(os.path.abspath(__file__)).get_logger()
 
 
@@ -39,13 +44,6 @@ def get_current_manager():
     return __DAG_CONTEXT__.get(None)
 
 
-def next_batch_job_id():
-    i = 0
-    while True:
-        yield i
-        i += 1
-
-
 class CyclicDependencyError(BatchJobError):
     """
     Exception raised when a cyclic dependency is detected in the batch job dependencies.
@@ -79,6 +77,12 @@ class MissingDagError(BatchJobError):
         )
 
 
+class CodePackageMismatchError(BatchJobError):
+    def __init__(self, message: str):
+        super().__init__(message)
+        self.message = message
+
+
 class BatchJob:
     """
     Represents a job within a batch job.
@@ -93,6 +97,7 @@ class BatchJob:
         resources: Optional[Resources] = None,
         dependencies: Sequence[Union[str, BatchJob]] = [],
         dag: Optional[Batch] = get_current_manager(),
+        argument_mapping: Dict[str, str] = {},
     ):
         self.name = name
         self.entrypoint = entrypoint
@@ -102,14 +107,46 @@ class BatchJob:
             raise MissingDagError()
         self.__id = dag.next_job_id
         self.dag = dag
+        self.__args: Dict[str, Any] = {}
         self.dag.__BATCH_JOB_NAMES__[self.name] = self.__id
         self.dependencies = self.translate_dependencies()
         assert all(
             isinstance(dep, int) for dep in self.dependencies
         ), "All dependencies must be integers representing job IDs."
         self.dag.add_job(self)
+        self.__argument_mapping = argument_mapping or {}
+
+    def __call__(self, *args, **kwds) -> BatchJob:
+        """
+        Allows the BatchJob instance to be called like a function, returning itself.
+        This is useful for chaining or functional-style programming.
+        """
+        return self
+
+    @property
+    def args(self) -> Dict[str, Any]:
+        """
+        Returns the arguments for the BatchJob instance.
+        """
+        return self.__args or {}
+
+    @args.setter
+    def args(self, args: Dict[str, Any]):
+        """
+        Sets the arguments for the BatchJob instance.
+        """
+        if not isinstance(args, dict):
+            raise TypeError(f"Expected a dictionary for args, got {type(args)}")
+        self.__args = args
+
+    @property
+    def id(self) -> int:
+        """
+        Returns the unique identifier of the BatchJob instance.
+        """
+        return self.__id
 
-    def alias(self, name: str):
+    def alias(self, name: str) -> BatchJob:
         """
         Set an alias for the BatchJob instance.
 
@@ -122,19 +159,48 @@ class BatchJob:
         self.name = name
         return self
 
+    def set_resources(
+        self,
+        resources: Optional[Resources] = None,
+        memory: Optional[str] = None,
+        cpu: Optional[float] = None,
+    ) -> BatchJob:
+        """
+        Set the resources for the BatchJob instance.
+
+        :param resources: The Resources instance to set.
+        """
+        if resources is not None:
+            if not isinstance(resources, Resources):
+                raise TypeError(f"Expected Resources instance, got {type(resources)}")
+        else:
+            resources = Resources(
+                memory=memory or DEFAULT_TASK_MEMORY, cpu=cpu or DEFAULT_TASK_CPU
+            )
+        self.resources = resources
+        return self
+
     def __repr__(self):
         return (
             f"BatchJob(name={self.name}, entrypoint={self.entrypoint}, "
             f"resources={self.resources}) (id={self.__id})"
         )
 
+    def __getstate__(self) -> object:
+        state = self.__dict__.copy()
+        state.pop("dag", None)
+        return state
+
+    def __setstate__(self, state: dict):
+        self.__dict__.update(state)
+
     def to_dict(self):
         """
         Convert the BatchJob instance to a dictionary representation.
         """
         return {
             "display_name": self.name,
-            "name": self.__id,
+            "child_number": self.__id,
             "entrypoint": str(self.entrypoint),
             "memory": self.resources.memory if self.resources else DEFAULT_TASK_MEMORY,
             "cpu": self.resources.cpu if self.resources else DEFAULT_TASK_CPU,
@@ -169,6 +235,9 @@ class BatchJob:
 
     def __add_dependency__(self, other):
         self.dependencies.add(other.__id)
+        arg_name = self.__argument_mapping.get(other.name, other.name)
+        if arg_name is not None:
+            self.__args[arg_name] = other
 
     def __lshift__(
         self, other: Sequence[BatchJob] | BatchJob
@@ -209,7 +278,13 @@ class BatchJob:
         Execute the job's entrypoint.
         """
         if isinstance(self.entrypoint, EntryPoint):
-            self.entrypoint()
+            env = {
+                "DATATAILR_BATCH_ID": str(self.dag.id),
+                "DATATAILR_JOB_ID": str(self.__id),
+                "DATATAILR_JOB_NAME": self.name,
+                "DATATAILR_JOB_ARGUMENT_MAPPING": encode_json(self.__argument_mapping),
+            }
+            self.entrypoint(env=env)
         else:
             raise TypeError(f"Invalid entrypoint type: {type(self.entrypoint)}")
 
@@ -223,12 +298,17 @@ class Batch(Job):
 
     def __init__(
         self,
-        environment: Optional[Environment],
         name: str,
-        image: Image,
-        run_as: Optional[Union[str, User]],
+        environment: Optional[Environment] = Environment.DEV,
+        schedule: Optional[Schedule] = None,
+        image: Optional[Image] = None,
+        run_as: Optional[Union[str, User]] = User.signed_user(),
         resources: Resources = Resources(memory="100m", cpu=1),
         acl: Optional[ACL] = None,
+        local_run: bool = False,
+        python_requirements: str = "",
+        build_script_pre: str = "",
+        build_script_post: str = "",
     ):
         super().__init__(
             environment=environment,
@@ -237,19 +317,25 @@ class Batch(Job):
             run_as=run_as,
             resources=resources,
             acl=acl,
+            python_requirements=python_requirements,
+            build_script_pre=build_script_pre,
+            build_script_post=build_script_post,
+            type=JobType.BATCH,
         )
-        self.type = JobType.BATCH
         self.__jobs: List[BatchJob] = []
         self._auto_run = False
-        self.__next_job_id = next_batch_job_id()
+        self.__next_job_id = -1
         self.__BATCH_JOB_NAMES__: Dict[str, int] = {}
+        self.__local_run = local_run
+        self.__schedule = schedule
 
     @property
     def next_job_id(self):
         """
         Returns a generator for the next job ID in the batch.
         """
-        return next(self.__next_job_id)
+        self.__next_job_id += 1
+        return self.__next_job_id
 
     def add_job(self, job: BatchJob):
         """
@@ -265,6 +351,25 @@ class Batch(Job):
             raise DuplicateJobNameError(job.name)
         # Use the batch level resource values as defaults for jobs
         job.resources = job.resources or self.resources
+        image_path_to_repo = self.image.path_to_repo
+        image_path_to_module = self.image.path_to_module
+        package_path_to_repo = job.entrypoint.path_to_repo
+        package_path_to_module = job.entrypoint.path_to_module
+
+        if image_path_to_repo is None:
+            self.image.path_to_repo = package_path_to_repo
+        elif package_path_to_repo != image_path_to_repo:
+            raise CodePackageMismatchError(
+                f"Function {job.entrypoint.function_name} is defined in a different package root: "
+                f"{package_path_to_repo} != {image_path_to_repo}"
+            )
+        if image_path_to_module is None:
+            self.image.path_to_module = package_path_to_module
+        elif package_path_to_module != image_path_to_module:
+            raise CodePackageMismatchError(
+                f"Function {job.entrypoint.function_name} is defined in a different module: "
+                f"{package_path_to_module} != {image_path_to_module}"
+            )
         self.__jobs.append(job)
 
     def is_job_in(self, job: BatchJob) -> bool:
@@ -279,6 +384,7 @@ class Batch(Job):
         """
         batch_dict = super().to_dict()
         batch_dict["jobs"] = [job.to_dict() for job in self.__jobs]
+        batch_dict["schedule"] = str(self.__schedule) if self.__schedule else None
         return batch_dict
 
     def to_json(self):
@@ -331,11 +437,41 @@ class Batch(Job):
                 "A cyclic dependency exists amongst {}".format(jobs)
             )
 
+    def get_schedule_args(self) -> Dict[str, Any]:
+        if isinstance(self.__schedule, Schedule):
+            args = {
+                "at_minute": self.__schedule.at_minutes,
+                "every_minute": self.__schedule.every_minute,
+                "at_hour": self.__schedule.at_hours,
+                "every_hour": self.__schedule.every_hour,
+                "weekdays": self.__schedule.weekdays,
+                "day_of_month": self.__schedule.day_of_month,
+                "in_month": self.__schedule.in_month,
+                "every_month": self.__schedule.every_month,
+                "timezone": self.__schedule.timezone,
+                "run_after_job_uuid": self.__schedule.run_after_job_uuid,
+                "run_after_job_name": self.__schedule.run_after_job_name,
+                "run_after_job_condition": self.__schedule.run_after_job_condition,
+            }
+            args = {key: value for key, value in args.items() if value is not None}
+            for key, value in args.items():
+                if isinstance(value, list):
+                    args[key] = ",".join(map(str, value))
+            return args
+        return {}
+
     def run(self) -> Tuple[bool, str]:
-        if is_dt_installed():
+        def arg_name(arg: Union[BatchJob, str]) -> str:
+            return arg.name if isinstance(arg, BatchJob) else arg
+
+        args = {
+            j.name: {k: arg_name(v) for k, v in j.args.items()} for j in self.__jobs
+        }
+        __ARGUMENTS_CACHE__.add_arguments(self.id, args)
+        if not self.__local_run and is_dt_installed():
             return super().run()
         else:
-            os.environ["DATATAILR_BATCH_RUN_ID"] = "1"
+            os.environ["DATATAILR_BATCH_RUN_ID"] = uuid.uuid4().hex[:8]
             for step in self.__topological_sort__():
                 for job_id in step:
                     job = self.__jobs[job_id]