avtomatika 1.0b10__py3-none-any.whl → 1.0b11__py3-none-any.whl

avtomatika/blueprint.py

@@ -131,7 +131,14 @@ class StateMachineBlueprint:
         self.name = name
         self.api_endpoint = api_endpoint
         self.api_version = api_version
-        self.data_stores: dict[str, AsyncDictStore] = data_stores if data_stores is not None else {}
+        self.data_stores: dict[str, AsyncDictStore] = {}
+        if data_stores:
+            for ds_name, ds_data in data_stores.items():
+                if isinstance(ds_data, AsyncDictStore):
+                    self.data_stores[ds_name] = ds_data
+                else:
+                    self.data_stores[ds_name] = AsyncDictStore(ds_data)
+
         self.handlers: dict[str, Callable] = {}
         self.aggregator_handlers: dict[str, Callable] = {}
         self.conditional_handlers: list[ConditionalHandler] = []

avtomatika/executor.py

@@ -238,6 +238,9 @@ class JobExecutor:
                             action_factory.sub_blueprint_to_run,
                             duration_ms,
                         )
+                    elif job_state["current_state"] in blueprint.end_states:
+                        status = JOB_STATUS_FINISHED if job_state["current_state"] == "finished" else JOB_STATUS_FAILED
+                        await self._handle_terminal_reached(job_state, status, duration_ms)
 
                 except Exception as e:
                     # This catches errors within the handler's execution.
@@ -248,6 +251,40 @@ class JobExecutor:
             if message_id in self._processing_messages:
                 self._processing_messages.remove(message_id)
 
+    async def _handle_terminal_reached(
+        self,
+        job_state: dict[str, Any],
+        status: str,
+        duration_ms: int,
+    ) -> None:
+        job_id = job_state["id"]
+        current_state = job_state["current_state"]
+        logger.info(f"Job {job_id} reached terminal state '{current_state}' with status '{status}'")
+
+        await self.history_storage.log_job_event(
+            {
+                "job_id": job_id,
+                "state": current_state,
+                "event_type": "job_completed",
+                "duration_ms": duration_ms,
+                "context_snapshot": job_state,
+            },
+        )
+
+        job_state["status"] = status
+        await self.storage.save_job_state(job_id, job_state)
+
+        # Clean up S3 files if service is available
+        s3_service = self.engine.app.get(S3_SERVICE_KEY)
+        if s3_service:
+            task_files = s3_service.get_task_files(job_id)
+            if task_files:
+                create_task(task_files.cleanup())
+
+        await self._check_and_resume_parent(job_state)
+        event_type = "job_finished" if status == JOB_STATUS_FINISHED else "job_failed"
+        await self.engine.send_job_webhook(job_state, event_type)
+
     async def _handle_transition(
         self,
         job_state: dict[str, Any],
@@ -270,28 +307,11 @@ class JobExecutor:
             },
         )
 
-        # When transitioning to a new state, reset the retry counter.
         job_state["retry_count"] = 0
         job_state["current_state"] = next_state
         job_state["status"] = JOB_STATUS_RUNNING
         await self.storage.save_job_state(job_id, job_state)
-
-        if next_state not in TERMINAL_STATES:
-            await self.storage.enqueue_job(job_id)
-        else:
-            logger.info(f"Job {job_id} reached terminal state {next_state}")
-
-            # Clean up S3 files if service is available
-            s3_service = self.engine.app.get(S3_SERVICE_KEY)
-            if s3_service:
-                task_files = s3_service.get_task_files(job_id)
-                if task_files:
-                    # Run cleanup in background to not block response
-                    create_task(task_files.cleanup())
-
-            await self._check_and_resume_parent(job_state)
-            event_type = "job_finished" if next_state == JOB_STATUS_FINISHED else "job_failed"
-            await self.engine.send_job_webhook(job_state, event_type)
+        await self.storage.enqueue_job(job_id)
 
     async def _handle_dispatch(
         self,

avtomatika/s3.py

@@ -335,12 +335,11 @@ class S3Service:
         try:
             self._store = S3Store(
                 bucket=self.config.S3_DEFAULT_BUCKET,
-                access_key_id=self.config.S3_ACCESS_KEY,
-                secret_access_key=self.config.S3_SECRET_KEY,
+                aws_access_key_id=self.config.S3_ACCESS_KEY,
+                aws_secret_access_key=self.config.S3_SECRET_KEY,
                 region=self.config.S3_REGION,
                 endpoint=self.config.S3_ENDPOINT_URL,
                 allow_http="http://" in self.config.S3_ENDPOINT_URL,
-                force_path_style=True,
             )
             self._semaphore = Semaphore(self.config.S3_MAX_CONCURRENCY)
             logger.info(

avtomatika/services/worker_service.py

@@ -104,7 +104,6 @@ class WorkerService:
 
         job_id = result_payload.get("job_id")
         task_id = result_payload.get("task_id")
-        result_data = result_payload.get("result", {})
 
         if not job_id or not task_id:
             raise ValueError("job_id and task_id are required")
@@ -113,25 +112,33 @@ class WorkerService:
         if not job_state:
             raise LookupError("Job not found")
 
+        result_status = result_payload.get("status", TASK_STATUS_SUCCESS)
+        worker_data_content = result_payload.get("data")
+
         if job_state.get("status") == JOB_STATUS_WAITING_FOR_PARALLEL:
             await self.storage.remove_job_from_watch(f"{job_id}:{task_id}")
-            job_state.setdefault("aggregation_results", {})[task_id] = result_data
 
-            branches = job_state.setdefault("active_branches", [])
-            if task_id in branches:
-                branches.remove(task_id)
+            def _update_parallel_results(state: dict[str, Any]) -> dict[str, Any]:
+                state.setdefault("aggregation_results", {})[task_id] = result_payload
+                branches = state.setdefault("active_branches", [])
+                if task_id in branches:
+                    branches.remove(task_id)
+
+                if not branches:
+                    state["status"] = JOB_STATUS_RUNNING
+                    state["current_state"] = state["aggregation_target"]
+                return state
 
-            if not branches:
+            updated_job_state = await self.storage.update_job_state_atomic(job_id, _update_parallel_results)
+
+            if not updated_job_state.get("active_branches"):
                 logger.info(f"All parallel branches for job {job_id} have completed.")
-                job_state["status"] = JOB_STATUS_RUNNING
-                job_state["current_state"] = job_state["aggregation_target"]
-                await self.storage.save_job_state(job_id, job_state)
                 await self.storage.enqueue_job(job_id)
             else:
+                remaining = len(updated_job_state["active_branches"])
                 logger.info(
-                    f"Branch {task_id} for job {job_id} completed. Waiting for {len(branches)} more.",
+                    f"Branch {task_id} for job {job_id} completed. Waiting for {remaining} more.",
                 )
-                await self.storage.save_job_state(job_id, job_state)
 
             return "parallel_branch_result_accepted"
 
@@ -148,14 +155,12 @@ class WorkerService:
                 "event_type": "task_finished",
                 "duration_ms": duration_ms,
                 "worker_id": authenticated_worker_id,
-                "context_snapshot": {**job_state, "result": result_data},
+                "context_snapshot": {**job_state, "result": result_payload},
             },
         )
 
-        result_status = result_data.get("status", TASK_STATUS_SUCCESS)  # Default to success? Constant?
-
         if result_status == TASK_STATUS_FAILURE:
-            return await self._handle_task_failure(job_state, task_id, result_data)
+            return await self._handle_task_failure(job_state, task_id, result_payload)
 
         if result_status == TASK_STATUS_CANCELLED:
             logger.info(f"Task {task_id} for job {job_id} was cancelled by worker.")
@@ -171,13 +176,11 @@ class WorkerService:
             return "result_accepted_cancelled"
 
         transitions = job_state.get("current_task_transitions", {})
-        result_status = result_data.get("status", TASK_STATUS_SUCCESS)
         next_state = transitions.get(result_status)
 
         if next_state:
             logger.info(f"Job {job_id} transitioning based on worker status '{result_status}' to state '{next_state}'")
 
-            worker_data_content = result_data.get("data")
             if worker_data_content and isinstance(worker_data_content, dict):
                 if "state_history" not in job_state:
                     job_state["state_history"] = {}
@@ -202,8 +205,8 @@ class WorkerService:
             await self.storage.save_job_state(job_id, job_state)
             return "result_accepted_failure"
 
-    async def _handle_task_failure(self, job_state: dict, task_id: str, result_data: dict) -> str:
-        error_details = result_data.get("error", {})
+    async def _handle_task_failure(self, job_state: dict, task_id: str, result_payload: dict) -> str:
+        error_details = result_payload.get("error", {})
         error_type = ERROR_CODE_TRANSIENT
         error_message = "No error details provided."
 

avtomatika/storage/base.py

@@ -90,6 +90,20 @@ class StorageBackend(ABC):
         """
         raise NotImplementedError
 
+    @abstractmethod
+    async def update_job_state_atomic(
+        self,
+        job_id: str,
+        update_callback: Any,
+    ) -> dict[str, Any]:
+        """Atomically update the state of a job using a callback function.
+
+        :param job_id: Unique identifier for the job.
+        :param update_callback: A callable that takes the current state and returns the updated state.
+        :return: The updated full state of the job.
+        """
+        raise NotImplementedError
+
     @abstractmethod
     async def register_worker(
         self,

avtomatika/storage/memory.py

@@ -62,6 +62,17 @@ class MemoryStorage(StorageBackend):
             self._jobs[job_id].update(update_data)
             return self._jobs[job_id]
 
+    async def update_job_state_atomic(
+        self,
+        job_id: str,
+        update_callback: Any,
+    ) -> dict[str, Any]:
+        async with self._lock:
+            current_state = self._jobs.get(job_id, {})
+            updated_state = update_callback(current_state)
+            self._jobs[job_id] = updated_state
+            return updated_state
+
     async def register_worker(
         self,
         worker_id: str,

avtomatika/storage/redis.py

@@ -97,6 +97,19 @@ class RedisStorage(StorageBackend):
         update_data: dict[str, Any],
     ) -> dict[str, Any]:
         """Atomically update the job state in Redis using a transaction."""
+
+        def _merge(state: dict[str, Any]) -> dict[str, Any]:
+            state.update(update_data)
+            return state
+
+        return await self.update_job_state_atomic(job_id, _merge)
+
+    async def update_job_state_atomic(
+        self,
+        job_id: str,
+        update_callback: Any,
+    ) -> dict[str, Any]:
+        """Atomically update the job state in Redis using a transaction and callback."""
         key = self._get_key(job_id)
 
         async with self._redis.pipeline(transaction=True) as pipe:
@@ -105,12 +118,11 @@ class RedisStorage(StorageBackend):
                     await pipe.watch(key)
                     current_state_raw = await pipe.get(key)
                     current_state: dict[str, Any] = self._unpack(current_state_raw) if current_state_raw else {}
-                    current_state.update(update_data)
-
+                    updated_state = update_callback(current_state)
                     pipe.multi()
-                    pipe.set(key, self._pack(current_state))
+                    pipe.set(key, self._pack(updated_state))
                     await pipe.execute()
-                    return current_state
+                    return updated_state
                 except WatchError:
                     continue
 

{avtomatika-1.0b10.dist-info → avtomatika-1.0b11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: avtomatika
-Version: 1.0b10
+Version: 1.0b11
 Summary: A state-machine based orchestrator for long-running AI and other jobs.
 Author-email: Dmitrii Gagarin <madgagarin@gmail.com>
 Project-URL: Homepage, https://github.com/avtomatika-ai/avtomatika
@@ -58,7 +58,6 @@ Dynamic: license-file
 
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 [![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/release/python-3110/)
-[![Tests](https://github.com/avtomatika-ai/avtomatika/actions/workflows/ci.yml/badge.svg)](https://github.com/avtomatika-ai/avtomatika/actions/workflows/ci.yml)
 [![Code Style: Ruff](https://img.shields.io/badge/code%20style-ruff-000000.svg)](https://github.com/astral-sh/ruff)
 
 Avtomatika is a powerful, state-driven engine for managing complex asynchronous workflows in Python. It provides a robust framework for building scalable and resilient applications by separating process logic from execution logic.

{avtomatika-1.0b10.dist-info → avtomatika-1.0b11.dist-info}/RECORD

@@ -1,7 +1,7 @@
 avtomatika/__init__.py,sha256=D5r3L-H06uxsY_wgfh7u9YR29QvZMer1BlvzjW9Umfo,701
 avtomatika/api.html,sha256=6Sj0vwAUZsbLKwlB58ONAttCB52e8h3fidspLOwMMGE,32894
 avtomatika/app_keys.py,sha256=Zd2TaGPduzyEFJgdPvgSH1skdBx2mX-Prj1ma9fAXRo,1275
-avtomatika/blueprint.py,sha256=ZRMis9LOtBwZ9MMHqMF7WAgfMxE0M6-xP5s1IwLgpow,11875
+avtomatika/blueprint.py,sha256=OPJShSdh8asl9G2kWzbFu1CKMzsq15fo37I0eYlISkg,12119
 avtomatika/client_config_loader.py,sha256=zVVHZlxSqZUaNpZ4zoU0T1CFYXdxy-3vKSmPcaFuHSY,2772
 avtomatika/compression.py,sha256=bhA1kw4YrCR3I3kdquZSY0fAzCrRrjtz55uepzLUDKI,2498
 avtomatika/config.py,sha256=27ov8BNbiUpkZ1sjtx3pifRavwcxJ_zUgIdkL_pgqv8,3595
@@ -11,7 +11,7 @@ avtomatika/data_types.py,sha256=D_IUzMW8zMz-_MaqVp9MG53rG37Cb3McyRZuIXxvdlE,1108
 avtomatika/datastore.py,sha256=gJjhZ5kxjF8pmbbPQb_qu3HPUpfy2c6T75KZ-smb_zg,545
 avtomatika/dispatcher.py,sha256=5J5GBWFfaGCGXUkM-2fhMeg2n2nTO0BH3ffkzsnSsaE,8784
 avtomatika/engine.py,sha256=Hb6MLanMjx1GDAfkbNJU-K4RXMuPZQP7_HA_0VR8WMw,20916
-avtomatika/executor.py,sha256=X5AU7hWflH8rSYKxl_wh2RhdYhpyktynmK8mcfJgT-8,24218
+avtomatika/executor.py,sha256=bu8_Xmr_hRNsatAKUzypIFWOZQT2yE_gMU4XfGBt4u4,24923
 avtomatika/health_checker.py,sha256=jXYSH4BPeZ4LCxSZV4uXM4BZhGJYgpoAOWQXE8yojLo,2078
 avtomatika/logging_config.py,sha256=cVY8aOeaWncsvkN015WgC74NTF6r55-OA3E1ux8P824,3347
 avtomatika/metrics.py,sha256=tiksK1fFSOMlz8zFu6GT19JTduvxMTNlLu0QFrTHoQI,1866
@@ -19,7 +19,7 @@ avtomatika/py.typed,sha256=CT_L7gw2MLcQY-X0vs-xB5Vr0wzvGo7GuQYPI_qwJE8,65
 avtomatika/quota.py,sha256=DNcaL6k0J1REeP8sVqbY9FprY_3BSr2SxM2Vf4mEqdw,1612
 avtomatika/ratelimit.py,sha256=hFGW5oN9G6_W_jnHmopXW8bRjjzlvanY19MLghsNLE8,1306
 avtomatika/reputation.py,sha256=pK-x9FrPN2Oc2gtPa1AZJHlhvkd7xlRe4orxM2auJJc,3979
-avtomatika/s3.py,sha256=I0fDw5I44RJAqSv4tREvwHp2cxB0mGY_l2cVZWpe3As,14110
+avtomatika/s3.py,sha256=Byc5C_KTo0mOErQRlhDJNPZplyqrWxKe4GSeU99Zaqk,14079
 avtomatika/scheduler.py,sha256=F5Kv5Rx34nDd0mE5jxjwpjRg8duDZBEr91N5Y6CNR24,4231
 avtomatika/scheduler_config_loader.py,sha256=38x-4G4yRrhSrLdmZ4aTb7WggE-BcGblKZO7x97nW6Y,1352
 avtomatika/security.py,sha256=eENEUc0OsHm6wN2H-ckGmiaV9qrZSbYsHFCWyYb3aLs,3271
@@ -34,15 +34,15 @@ avtomatika/history/noop.py,sha256=hLzt0RblsrKUtoyQNauOni6jCi-IYCWEPsiR0vh7tho,12
 avtomatika/history/postgres.py,sha256=T0XpDurnh48pPI-2JhB285GdNIexNkCSu8ExhLJzcxc,9538
 avtomatika/history/sqlite.py,sha256=txWax9RVzBQzIZuU-SjHnEXEzBmGzIjqzoVsK2oyiAQ,9252
 avtomatika/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-avtomatika/services/worker_service.py,sha256=cPik-DUYPka0lOT38fvmVpEnmt1_n2p44z9emyikdLg,11518
+avtomatika/services/worker_service.py,sha256=lFHVqbXG-4v5Ec17FZCHaB6Uu8U2nWOGpPZOKjzwM00,11596
 avtomatika/storage/__init__.py,sha256=mGRj_40dWZ7R7uYbqC6gCsUWCKHAbZz4ZVIhYg5dT_E,262
-avtomatika/storage/base.py,sha256=Tb_4fF0Vr10cgoXepA-1YUSgi27qYKQ7Qz1Y87XiRII,13375
-avtomatika/storage/memory.py,sha256=23eNAcEleM6Yqi_kSn-dLEBJRMzrhlgRHVgxrVhZPrk,14560
-avtomatika/storage/redis.py,sha256=MPSQRuAzWNtKQZco_5ExvCpaWbs5_80or5QVrU1GcIM,20235
+avtomatika/storage/base.py,sha256=54II8RfrEQzCT9NH_ECorM9SdvM-e5f-_MRtJInBczw,13856
+avtomatika/storage/memory.py,sha256=DqUd7SQmneJCNd-YaWLQL-Gpz3FwRHFAH6xx2CIIqY4,14915
+avtomatika/storage/redis.py,sha256=fN0e3_2CP-8H1WKTBljXeHMVkrzLXs8znE2YkKbjuy0,20658
 avtomatika/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 avtomatika/utils/webhook_sender.py,sha256=LoJ6z_1p-OngjPYl9Pk1N1t9xrP6-v-7xOg_AmWPuVc,3644
-avtomatika-1.0b10.dist-info/licenses/LICENSE,sha256=tqCjw9Y1vbU-hLcWi__7wQstLbt2T1XWPdbQYqCxuWY,1072
-avtomatika-1.0b10.dist-info/METADATA,sha256=iFXea4IsLOM9dqBBcMDPqbEZDGTMbLcQoV-7SDpk_xQ,28593
-avtomatika-1.0b10.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-avtomatika-1.0b10.dist-info/top_level.txt,sha256=gLDWhA_wxHj0I6fG5X8vw9fE0HSN4hTE2dEJzeVS2x8,11
-avtomatika-1.0b10.dist-info/RECORD,,
+avtomatika-1.0b11.dist-info/licenses/LICENSE,sha256=tqCjw9Y1vbU-hLcWi__7wQstLbt2T1XWPdbQYqCxuWY,1072
+avtomatika-1.0b11.dist-info/METADATA,sha256=3Ok1N9NHujsXKORqv28ZSZmfHpIJ3_GjWGDkXbVyGSU,28432
+avtomatika-1.0b11.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+avtomatika-1.0b11.dist-info/top_level.txt,sha256=gLDWhA_wxHj0I6fG5X8vw9fE0HSN4hTE2dEJzeVS2x8,11
+avtomatika-1.0b11.dist-info/RECORD,,