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

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,12 +107,14 @@ 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:
         """
@@ -116,6 +123,22 @@ class BatchJob:
         """
         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:
         """
@@ -123,7 +146,7 @@ class BatchJob:
         """
         return self.__id
 
-    def alias(self, name: str):
+    def alias(self, name: str) -> BatchJob:
         """
         Set an alias for the BatchJob instance.
 
@@ -136,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,
@@ -183,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
@@ -223,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)}")
 
@@ -237,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,
@@ -251,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):
         """
@@ -279,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:
@@ -293,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):
@@ -345,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]

datatailr/scheduler/batch_decorator.py

@@ -11,21 +11,25 @@
 import functools
 import inspect
 import os
+from typing import Callable
 
 from datatailr.logging import DatatailrLogger
-from datatailr.scheduler.arguments_cache import ArgumentsCache
+from datatailr.scheduler.arguments_cache import ArgumentsCache, CacheNotFoundError
 from datatailr.scheduler.base import EntryPoint, JobType, Resources
-from datatailr.scheduler.batch import BatchJob, get_current_manager
+from datatailr.scheduler.batch import (
+    BatchJob,
+    get_current_manager,
+)
 from datatailr.scheduler.constants import DEFAULT_TASK_CPU, DEFAULT_TASK_MEMORY
 from datatailr.scheduler.utils import get_available_env_args
 
 __ARGUMENTS_CACHE__ = ArgumentsCache()
-__FUNCTIONS_CREATED_IN_DAG__: dict[BatchJob, str] = {}
+__FUNCTIONS_CREATED_IN_DAG__: dict[Callable, str] = {}
 logger = DatatailrLogger(__name__).get_logger()
 
 
 def batch_run_id() -> str:
-    return os.environ["DATATAILR_BATCH_RUN_ID"]
+    return os.getenv("DATATAILR_BATCH_RUN_ID", "unknown")
 
 
 def dag_id(job: BatchJob) -> str:
@@ -34,7 +38,7 @@ def dag_id(job: BatchJob) -> str:
     )
 
 
-def batch_decorator(memory=DEFAULT_TASK_MEMORY, cpu=DEFAULT_TASK_CPU):
+def batch_decorator(memory: str = DEFAULT_TASK_MEMORY, cpu: float = DEFAULT_TASK_CPU):
     """
     Decorator to mark a function as a batch job.
     This decorator can be used to wrap functions that should be executed as part of batch jobs.
@@ -55,39 +59,62 @@ def batch_decorator(memory=DEFAULT_TASK_MEMORY, cpu=DEFAULT_TASK_CPU):
                 # There are two possible scenarios:
                 # 1. The function is called directly, not as part of a batch job. In this case, the args and kwargs should be used.
                 # 2. The function is called as part of a batch job - it was constructed as part of a DAG and is now being executed.
-                if func not in __FUNCTIONS_CREATED_IN_DAG__:
-                    return func(*args, **kwargs)
-                function_arguments = [v.name for v in parameters.values()]
                 env_args = get_available_env_args()
+                all_function_args = [
+                    v.name
+                    for v in parameters.values()
+                    if v.kind
+                    not in (
+                        inspect.Parameter.VAR_POSITIONAL,
+                        inspect.Parameter.VAR_KEYWORD,
+                    )
+                ]
                 final_args = list(args)
-                final_kwargs = kwargs.copy()
 
                 for name, value in env_args.items():
-                    if name in function_arguments:
-                        if len(final_args) < len(function_arguments):
+                    if name in all_function_args:
+                        if len(final_args) < len(all_function_args):
                             final_args.extend(
-                                [None] * (len(function_arguments) - len(final_args))
+                                [None] * (len(all_function_args) - len(final_args))
                             )
-                        final_args[function_arguments.index(name)] = value
-                function_arguments = __ARGUMENTS_CACHE__.get_arguments(
-                    dag_id(func), func.__name__
+                        final_args[all_function_args.index(name)] = value
+                try:
+                    final_kwargs = __ARGUMENTS_CACHE__.get_arguments(
+                        dag_id(func),
+                        os.getenv("DATATAILR_JOB_NAME", func.__name__),
+                        os.getenv("DATATAILR_BATCH_RUN_ID"),
+                    )
+                except CacheNotFoundError:
+                    final_kwargs = kwargs
+
+                if varargs is not None and varkw is None:
+                    for key in list(final_kwargs.keys()):
+                        if key not in parameters:
+                            final_args.append(final_kwargs.pop(key))
+
+                # Some of the loaded arguments are actually args and not kwargs.
+                if len(final_args) == len(parameters.keys()):
+                    for i, arg_name in enumerate(parameters.keys()):
+                        final_args[i] = final_kwargs.pop(arg_name, final_args[i])
+                result = func(*final_args, **final_kwargs)
+                __ARGUMENTS_CACHE__.add_result(
+                    batch_run_id(),
+                    os.getenv("DATATAILR_JOB_NAME", func.__name__),
+                    result,
                 )
-                result = func(**function_arguments)
-                __ARGUMENTS_CACHE__.add_result(batch_run_id(), func.__name__, result)
                 return result
             else:
-                __FUNCTIONS_CREATED_IN_DAG__[func] = dag.id
-                all_args = dict(zip(spec.args, args)) | kwargs
-
-                __ARGUMENTS_CACHE__.add_arguments(dag.id, func.__name__, all_args)
-
+                if varargs is not None:
+                    all_args = {job.name: job for job in args}
+                else:
+                    all_args = dict(zip(spec.args, args)) | kwargs
                 dag.set_autorun(True)
+
                 job = BatchJob(
                     name=func.__name__,
                     entrypoint=EntryPoint(
                         JobType.BATCH,
-                        module_name=func.__module__,
-                        function_name=func.__name__,
+                        func=func,
                     ),
                     resources=Resources(memory=memory, cpu=cpu),
                     dependencies=[
@@ -97,6 +124,8 @@ def batch_decorator(memory=DEFAULT_TASK_MEMORY, cpu=DEFAULT_TASK_CPU):
                     ],
                     dag=dag,
                 )
+                job.args = all_args
+                __FUNCTIONS_CREATED_IN_DAG__[job.entrypoint.func] = dag.id
                 return job
 
         module = inspect.getmodule(func)

datatailr/scheduler/constants.py

@@ -9,7 +9,7 @@
 # *************************************************************************
 
 DEFAULT_TASK_MEMORY = "100m"
-DEFAULT_TASK_CPU = 1
+DEFAULT_TASK_CPU = 1.0
 
 BATCH_JOB_ARGUMENTS = (
     "rundate",

datatailr/scheduler/schedule.py

@@ -0,0 +1,117 @@
+##########################################################################
+#
+#  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.
+##########################################################################
+
+from typing import Any
+from datatailr.wrapper import dt__Job
+import re
+
+
+__CLIENT__ = dt__Job()
+
+
+class Schedule:
+    """
+    Represents a schedule for batch jobs.
+    """
+
+    def __init__(
+        self,
+        cron_expression: str = "",
+        at_minutes: list[int] | None = None,
+        every_minute: int | None = None,
+        at_hours: list[int] | None = None,
+        every_hour: int | None = None,
+        weekdays: list[str] | None = None,
+        day_of_month: int | None = None,
+        in_month: list[str] | None = None,
+        every_month: int | None = None,
+        timezone: str | None = None,
+        run_after_job_uuid: str | None = None,
+        run_after_job_name: str | None = None,
+        run_after_job_condition: str | None = None,
+    ):
+        self.at_minutes = at_minutes
+        self.every_minute = every_minute
+        self.at_hours = at_hours
+        self.every_hour = every_hour
+        self.weekdays = weekdays
+        self.day_of_month = day_of_month
+        self.in_month = in_month
+        self.every_month = every_month
+        self.timezone = timezone
+        self.run_after_job_uuid = run_after_job_uuid
+        self.run_after_job_name = run_after_job_name
+        self.run_after_job_condition = run_after_job_condition
+        self.schedule_expression = None
+        self.cron_expression = cron_expression
+
+        self.__is_set__ = False
+
+    def __str__(self) -> str:
+        self.__compile__()
+        return self.cron_expression
+
+    def __repr__(self) -> str:
+        self.__compile__()
+        return f"Schedule(cron_expression={self.cron_expression}, timezone={self.timezone}) - {self.schedule_expression}"
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        super().__setattr__(name, value)
+        if name in [
+            "at_minutes",
+            "at_hours",
+            "every_minute",
+            "every_hour",
+            "weekdays",
+            "day_of_month",
+            "in_month",
+            "every_month",
+        ]:
+            self.__is_set__ = False
+
+    def __compile__(self):
+        if self.__is_set__:
+            return
+        argument_name = [
+            "at_minutes",
+            "at_hours",
+            "every_minute",
+            "every_hour",
+            "weekdays",
+            "day_of_month",
+            "in_month",
+            "every_month",
+            "run_after_job_uuid",
+            "run_after_job_name",
+            "run_after_job_condition",
+        ]
+        arguments = {}
+
+        for key in argument_name:
+            if hasattr(self, key) and getattr(self, key) is not None:
+                value = getattr(self, key)
+                if isinstance(value, list):
+                    value = ",".join(map(str, value))
+                arguments[key] = value
+
+        result = __CLIENT__.run("", cron_string=True, **arguments)
+        match = re.match(r"^(.*?)\s*\((.*?)\)$", result)
+        if match:
+            cron_expression, schedule_expression = match.groups()
+            self.cron_expression = cron_expression.strip()
+            self.schedule_expression = schedule_expression.strip()
+        self.__is_set__ = True
+
+    def get_cron_string(self) -> str:
+        """
+        Returns the compiled cron string.
+        """
+        self.__compile__()
+        return self.cron_expression

datatailr/scheduler/utils.py

@@ -16,7 +16,9 @@ from datatailr.scheduler.constants import BATCH_JOB_ARGUMENTS
 def get_available_env_args():
     """
     Get the available environment variables for batch job arguments.
-    This function retrieves the environment variables that match the keys defined in BATCH_JOB_ARGUMENTS.
+
+    This function retrieves the environment variables that match the keys defined in DATATAILR_BATCH_JOB_ARGUMENTS.
+
     Returns:
         dict: A dictionary of available environment variables for batch jobs.
     """