datatailr 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

datatailr/scheduler/__init__.py

@@ -0,0 +1,38 @@
+# *************************************************************************
+#
+#  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 datatailr.errors import BatchJobError, DatatailrError
+from datatailr.scheduler.base import (
+    ACL,
+    EntryPoint,
+    Environment,
+    Job,
+    JobType,
+    Resources,
+    User,
+)
+from datatailr.scheduler.batch import Batch, BatchJob, DuplicateJobNameError
+from datatailr.scheduler.batch_decorator import batch_decorator as batch
+
+__all__ = [
+    "Job",
+    "JobType",
+    "Environment",
+    "User",
+    "Resources",
+    "ACL",
+    "EntryPoint",
+    "Batch",
+    "BatchJob",
+    "batch",
+    "DatatailrError",
+    "BatchJobError",
+    "DuplicateJobNameError",
+]

datatailr/scheduler/arguments_cache.py

@@ -0,0 +1,126 @@
+##########################################################################
+#
+#  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.
+##########################################################################
+
+"""
+Module for caching arguments passed to batch jobs.
+
+This module provides two backends for caching:
+1. In-memory cache for local runs.
+2. Persistent cache using the dt__Blob module for remote runs.
+
+The cache stores arguments as a dictionary of dictionaries, where the outer dictionary's keys are job names
+and the inner dictionaries contain the arguments.
+
+This module is for internal use of the datatailr package.
+"""
+
+from collections import defaultdict
+from typing import Any, Dict
+
+from datatailr import is_dt_installed
+from datatailr.scheduler import BatchJob
+
+
+class ArgumentsCache:
+    def __init__(self, use_persistent_cache: bool = is_dt_installed()):
+        """
+        Initialize the 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)
+        )
+
+    def add_arguments(self, batch_run_id: str, job: str, arguments: Dict[str, Any]):
+        """
+        Add arguments to the cache for a specific job and batch run.
+
+        :param batch_run_id: Identifier for the batch run.
+        :param job_name: Name of the job.
+        :param arguments: Dictionary of arguments to store.
+        """
+        if self.use_persistent_cache and isinstance(job, str):
+            self._add_to_persistent_cache(batch_run_id, job, arguments)
+        else:
+            self.in_memory_cache[batch_run_id][job]["args"] = arguments
+
+    def get_arguments(self, batch_run_id: str, job: str) -> Dict[str, Any]:
+        """
+        Retrieve arguments from the cache for a specific job and batch run.
+
+        :param batch_run_id: Identifier for the batch run.
+        :param job_name: Name of the job.
+        :return: Dictionary of arguments.
+        """
+        if self.use_persistent_cache and isinstance(job, str):
+            return self._get_from_persistent_cache(batch_run_id, job)
+        arguments = {}
+        for key, value in (
+            self.in_memory_cache.get(batch_run_id, {})
+            .get(job, {})
+            .get("args", {})
+            .items()
+        ):
+            if isinstance(value, BatchJob):
+                arguments[key] = value.name
+            else:
+                arguments[key] = value
+        return arguments
+
+    def add_result(self, batch_run_id: str, job: str, result: Any):
+        """
+        Add the result of a batch job to the cache.
+
+        :param batch_run_id: Identifier for the batch run.
+        :param job: Name of the job.
+        :param result: Result of the batch job.
+        """
+        if self.use_persistent_cache and isinstance(job, str):
+            self._add_to_persistent_cache(batch_run_id, job, {"result": result})
+        else:
+            self.in_memory_cache[batch_run_id][job]["result"] = result
+
+    def get_result(self, batch_run_id: str, job: str) -> Any:
+        """
+        Retrieve the result of a batch job from the cache.
+
+        :param batch_run_id: Identifier for the batch run.
+        :param job: Name of the job.
+        :return: Result of the batch job.
+        """
+        if self.use_persistent_cache and isinstance(job, str):
+            return self._get_from_persistent_cache(batch_run_id, job).get("result")
+        return self.in_memory_cache[batch_run_id][job].get("result")
+
+    def _add_to_persistent_cache(
+        self, batch_run_id: str, job_name: str, arguments: Dict[str, Any]
+    ):
+        """
+        Add arguments to the persistent cache.
+
+        :param batch_run_id: Identifier for the batch run.
+        :param job_name: Name of the job.
+        :param arguments: Dictionary of arguments to store.
+        """
+        pass
+
+    def _get_from_persistent_cache(
+        self, batch_run_id: str, job_name: str
+    ) -> Dict[str, Any]:
+        """
+        Retrieve arguments from the persistent cache.
+
+        :param batch_run_id: Identifier for the batch run.
+        :param job_name: Name of the job.
+        :return: Dictionary of arguments.
+        """
+        return {}

datatailr/scheduler/base.py

@@ -0,0 +1,238 @@
+# *************************************************************************
+#
+#  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 importlib
+import json
+import os
+import subprocess
+import tempfile
+import uuid
+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.build.image import Image
+from datatailr.logging import DatatailrLogger
+
+logger = DatatailrLogger(os.path.abspath(__file__)).get_logger()
+
+
+class JobType(Enum):
+    """
+    Enum representing different types of DataTailr jobs.
+    """
+
+    BATCH = "batch"
+    SERVICE = "service"
+    APP = "app"
+    UNKNOWN = "unknown"
+
+    def __str__(self):
+        return self.value
+
+    def __repr__(self):
+        return f"JobType.{self.name}('{self.value}')"
+
+
+@dataclass
+class Resources:
+    """
+    Represents the resources required for a job.
+    """
+
+    memory: str = "100m"
+    cpu: int = 1
+
+
+class EntryPoint:
+    """
+    Represents an entry point for a DataTailr job.
+    This can be a function or a callable object.
+    """
+
+    def __init__(
+        self,
+        type: JobType,
+        func: Optional[Callable] = None,
+        module_name: Optional[str] = None,
+        function_name: Optional[str] = None,
+    ):
+        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.type = type
+
+    def __call__(self, *args, **kwargs):
+        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
+            return func(*args, **kwargs)
+
+        elif self.type == JobType.SERVICE:
+            raise NotImplementedError("Service jobs are not yet implemented.")
+
+        elif self.type == JobType.APP:
+            raise NotImplementedError("App jobs are not yet implemented.")
+
+    def __repr__(self):
+        return f"EntryPoint({self.function_name} from {self.module_name}, type={self.type})"
+
+    def __str__(self):
+        return f"{self.module_name}.{self.function_name}"
+
+
+class Job:
+    def __init__(
+        self,
+        environment: Optional[Environment],
+        name: str,
+        image: Image,
+        run_as: Optional[Union[str, User]],
+        resources: Resources = Resources(memory="100m", cpu=1),
+        acl: Optional[ACL] = None,
+    ):
+        if run_as is None:
+            run_as = User.signed_user()
+        if environment is None:
+            environment = Environment.DEV
+        elif isinstance(environment, str):
+            environment = Environment(environment.lower())
+        if isinstance(environment, str):
+            environment = Environment(environment)
+        self.acl = acl or ACL(user=run_as)
+        self.environment = environment
+        self.name = name
+        self.run_as = run_as
+        self.resources = resources
+        self.image = image
+
+        # Placeholders, to be set in derived classes
+        self.type: JobType = JobType.UNKNOWN
+        self.entrypoint = None
+        self.__id = str(uuid.uuid4())
+
+    @property
+    def id(self) -> str:
+        """
+        Unique identifier for the job.
+        """
+        return self.__id
+
+    def __repr__(self):
+        return (
+            f"Job(name={self.name}, environment={self.environment}, "
+            f"run_as={self.run_as}, resources={self.resources}, "
+            f"acl={self.acl}, type={self.type}, "
+            f"entrypoint={self.entrypoint}, image={self.image})"
+        )
+
+    def to_dict(self):
+        """
+        Convert the Job instance to a dictionary representation.
+        """
+        job_dict = {
+            "environment": str(self.environment),
+            "image": self.image.to_dict(),
+            "type": str(self.type) if self.type else None,
+            "name": self.name,
+            "run_as": self.run_as.name
+            if isinstance(self.run_as, User)
+            else self.run_as,
+            "acl": self.acl.to_dict(),
+        }
+        if self.type != JobType.BATCH:
+            job_dict["entrypoint"] = str(self.entrypoint) if self.entrypoint else None
+            job_dict["image"] = (self.image,)
+            job_dict["memory"] = (self.resources.memory,)
+            job_dict["cpu"] = self.resources.cpu
+        return job_dict
+
+    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
+        )
+        if not is_committed:
+            return (
+                False,
+                "Uncommitted changes detected. 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"),
+                shell=True,
+                capture_output=True,
+                text=True,
+            )
+            .stdout.strip()
+            .split("\t")[0]
+        )
+
+        if local_commit != remote_commit:
+            return (
+                False,
+                "Local commit does not match remote HEAD. Please pull the latest changes before running the job.",
+            )
+
+        branch = subprocess.run(
+            ("git rev-parse --abbrev-ref HEAD"),
+            shell=True,
+            capture_output=True,
+            text=True,
+        ).stdout.strip()
+        return True, ""
+
+    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]:
+                return False, check_result[1]
+            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())
+
+            dt__Job().run(f"file://{temp_file.name}")
+            os.remove(temp_file.name)
+
+            return True, ""
+        else:
+            raise NotImplementedError(
+                "DataTailr is not installed. Please install DataTailr to run this job."
+            )