1from __future__ import annotations 2 3from daggerml.contrib.executors._base import ExecutorBase 4from daggerml.contrib.executors.batch import BatchExecutor 5from daggerml.contrib.executors.docker import DockerExecutor 6from daggerml.contrib.executors.script import ScriptExecutor 7from daggerml.contrib.executors.ssh import SshExecutor 8 9__all__ = ["ExecutorBase", "BatchExecutor", "DockerExecutor", "ScriptExecutor", "SshExecutor"]
13class ExecutorBase: 14 """Base class for all executors. 15 16 The runtime owns durable adapter state. Executors receive ``adapter_state=None`` 17 on first launch and persisted state on later status checks. Executors 18 return terminal or in-progress result dicts via stdout/return value: 19 20 {"status": "retry", "error": None, "state": {...}} 21 {"status": "success", "error": None, "state": None} 22 {"status": "failure", "error": "<msg>", "state": None} 23 """ 24 25 name: str = "" 26 adapter: str = "" 27 28 # ------------------------------------------------------------------ 29 # Subclass interface 30 # ------------------------------------------------------------------ 31 32 def start(self, cache_key, execution_id, runnable, remote, scratch_uri) -> dict[str, Any]: 33 """Launch execution and return a result dict. 34 35 For synchronous executors this should return the terminal result 36 immediately. For async executors, return the durable resume state in the 37 initial ``retry`` result. 38 """ 39 raise NotImplementedError 40 41 def poll(self, cache_key, execution_id, runnable, state, remote, scratch_uri) -> dict[str, Any]: 42 """Check an in-flight job and return a result dict. 43 44 ``state`` is the immutable launch-time state returned by ``start()``. 45 Return a terminal result when done, or ``{"status": "retry", 46 "error": None, "state": ...}`` while still running. 47 Later returned state may be ignored by the runtime. 48 """ 49 raise NotImplementedError 50 51 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref) -> dict[str, Any]: 52 """Idempotently prune resources after a result was published.""" 53 del cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref 54 return {"status": "success", "error": None} 55 56 def cancel( 57 self, cache_key, execution_id, runnable, state, remote, scratch_uri, cancel_requested_by, argv_ref=None 58 ) -> dict[str, Any]: 59 raise NotImplementedError("This executor does not support cancellation") 60 61 # ------------------------------------------------------------------ 62 # Main dispatch 63 # ------------------------------------------------------------------ 64 65 @classmethod 66 def handle( 67 cls, 68 **payload: Any, 69 ) -> dict[str, Any]: 70 """Dispatch an explicit adapter operation to the executor.""" 71 operation = payload.get("operation") 72 if operation not in {"invoke", "cleanup", "cancel"}: 73 raise DmlRepoError(f"Unsupported adapter operation: {operation}") 74 required = { 75 "invoke": {"operation", "cache_key", "execution_id", "remote", "runnable", "adapter_state", "scratch_uri"}, 76 "cleanup": { 77 "operation", 78 "cache_key", 79 "execution_id", 80 "remote", 81 "runnable", 82 "adapter_state", 83 "scratch_uri", 84 "result_ref", 85 }, 86 "cancel": { 87 "operation", 88 "cache_key", 89 "execution_id", 90 "argv_ref", 91 "remote", 92 "runnable", 93 "adapter_state", 94 "scratch_uri", 95 "requested_by", 96 }, 97 }[operation] 98 if set(payload) != required: 99 raise DmlRepoError(f"Invalid {operation} adapter request fields") 100 cache_key = payload["cache_key"] 101 execution_id = payload["execution_id"] 102 remote = payload["remote"] 103 runnable = payload["runnable"] 104 adapter_state = payload["adapter_state"] 105 scratch_uri = payload["scratch_uri"] 106 if not all(isinstance(value, str) and value for value in (cache_key, execution_id, scratch_uri)): 107 raise DmlRepoError("Adapter request requires non-empty string identifiers") 108 if ( 109 not isinstance(remote, dict) 110 or set(remote) != {"root"} 111 or not isinstance(remote["root"], str) 112 or not remote["root"] 113 ): 114 raise DmlRepoError("Adapter request requires remote with non-empty root") 115 if not isinstance(runnable, dict): 116 raise DmlRepoError("Adapter request runnable must be an object") 117 if adapter_state is not None and not isinstance(adapter_state, dict): 118 raise DmlRepoError("adapter_state must be an object or null") 119 requested_by = payload.get("requested_by") 120 argv_ref = payload.get("argv_ref") 121 try: 122 valid_argv_ref = isinstance(argv_ref, str) and Ref(argv_ref).ns() == "node-argv" 123 except (TypeError, ValueError): 124 valid_argv_ref = False 125 if operation == "cancel" and not valid_argv_ref: 126 raise DmlRepoError("Cancel operation requires a node-argv ref") 127 if operation == "cancel" and requested_by is not None and not isinstance(requested_by, str): 128 raise DmlRepoError("Cancel operation requested_by must be a string or null") 129 try: 130 valid_result_ref = isinstance(payload.get("result_ref"), str) and Ref(payload["result_ref"]).ns() == "dag" 131 except (TypeError, ValueError): 132 valid_result_ref = False 133 if operation == "cleanup" and not valid_result_ref: 134 raise DmlRepoError("Cleanup operation requires a non-null result_ref") 135 executor = cls() 136 if operation == "cancel": 137 result = executor.cancel( 138 cache_key=cache_key, 139 execution_id=execution_id, 140 runnable=runnable, 141 state=adapter_state, 142 remote=remote, 143 scratch_uri=scratch_uri, 144 cancel_requested_by=requested_by, 145 argv_ref=argv_ref, 146 ) 147 elif operation == "cleanup": 148 result = executor.cleanup( 149 cache_key=cache_key, 150 execution_id=execution_id, 151 runnable=runnable, 152 state=adapter_state, 153 remote=remote, 154 scratch_uri=scratch_uri, 155 result_ref=payload["result_ref"], 156 ) 157 elif adapter_state is None: 158 result = executor.start( 159 cache_key=cache_key, 160 execution_id=execution_id, 161 runnable=runnable, 162 remote=remote, 163 scratch_uri=scratch_uri, 164 ) 165 else: 166 result = executor.poll( 167 cache_key=cache_key, 168 execution_id=execution_id, 169 runnable=runnable, 170 state=adapter_state, 171 remote=remote, 172 scratch_uri=scratch_uri, 173 ) 174 if not isinstance(result, dict): 175 raise DmlRepoError("Executor response must be a JSON object") 176 if "state" in result and "adapter_state" in result: 177 raise DmlRepoError("Executor response cannot contain both state and adapter_state") 178 if "state" in result: 179 result["adapter_state"] = result.pop("state") 180 return validate_adapter_response( 181 result, 182 success_status="cancelled" if operation == "cancel" else "success", 183 )
Base class for all executors.
The runtime owns durable adapter state. Executors receive adapter_state=None
on first launch and persisted state on later status checks. Executors
return terminal or in-progress result dicts via stdout/return value:
{"status": "retry", "error": None, "state": {...}}
{"status": "success", "error": None, "state": None}
{"status": "failure", "error": "<msg>", "state": None}
32 def start(self, cache_key, execution_id, runnable, remote, scratch_uri) -> dict[str, Any]: 33 """Launch execution and return a result dict. 34 35 For synchronous executors this should return the terminal result 36 immediately. For async executors, return the durable resume state in the 37 initial ``retry`` result. 38 """ 39 raise NotImplementedError
Launch execution and return a result dict.
For synchronous executors this should return the terminal result
immediately. For async executors, return the durable resume state in the
initial retry result.
41 def poll(self, cache_key, execution_id, runnable, state, remote, scratch_uri) -> dict[str, Any]: 42 """Check an in-flight job and return a result dict. 43 44 ``state`` is the immutable launch-time state returned by ``start()``. 45 Return a terminal result when done, or ``{"status": "retry", 46 "error": None, "state": ...}`` while still running. 47 Later returned state may be ignored by the runtime. 48 """ 49 raise NotImplementedError
Check an in-flight job and return a result dict.
state is the immutable launch-time state returned by start().
Return a terminal result when done, or {"status": "retry",
"error": None, "state": ...} while still running.
Later returned state may be ignored by the runtime.
51 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref) -> dict[str, Any]: 52 """Idempotently prune resources after a result was published.""" 53 del cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref 54 return {"status": "success", "error": None}
Idempotently prune resources after a result was published.
65 @classmethod 66 def handle( 67 cls, 68 **payload: Any, 69 ) -> dict[str, Any]: 70 """Dispatch an explicit adapter operation to the executor.""" 71 operation = payload.get("operation") 72 if operation not in {"invoke", "cleanup", "cancel"}: 73 raise DmlRepoError(f"Unsupported adapter operation: {operation}") 74 required = { 75 "invoke": {"operation", "cache_key", "execution_id", "remote", "runnable", "adapter_state", "scratch_uri"}, 76 "cleanup": { 77 "operation", 78 "cache_key", 79 "execution_id", 80 "remote", 81 "runnable", 82 "adapter_state", 83 "scratch_uri", 84 "result_ref", 85 }, 86 "cancel": { 87 "operation", 88 "cache_key", 89 "execution_id", 90 "argv_ref", 91 "remote", 92 "runnable", 93 "adapter_state", 94 "scratch_uri", 95 "requested_by", 96 }, 97 }[operation] 98 if set(payload) != required: 99 raise DmlRepoError(f"Invalid {operation} adapter request fields") 100 cache_key = payload["cache_key"] 101 execution_id = payload["execution_id"] 102 remote = payload["remote"] 103 runnable = payload["runnable"] 104 adapter_state = payload["adapter_state"] 105 scratch_uri = payload["scratch_uri"] 106 if not all(isinstance(value, str) and value for value in (cache_key, execution_id, scratch_uri)): 107 raise DmlRepoError("Adapter request requires non-empty string identifiers") 108 if ( 109 not isinstance(remote, dict) 110 or set(remote) != {"root"} 111 or not isinstance(remote["root"], str) 112 or not remote["root"] 113 ): 114 raise DmlRepoError("Adapter request requires remote with non-empty root") 115 if not isinstance(runnable, dict): 116 raise DmlRepoError("Adapter request runnable must be an object") 117 if adapter_state is not None and not isinstance(adapter_state, dict): 118 raise DmlRepoError("adapter_state must be an object or null") 119 requested_by = payload.get("requested_by") 120 argv_ref = payload.get("argv_ref") 121 try: 122 valid_argv_ref = isinstance(argv_ref, str) and Ref(argv_ref).ns() == "node-argv" 123 except (TypeError, ValueError): 124 valid_argv_ref = False 125 if operation == "cancel" and not valid_argv_ref: 126 raise DmlRepoError("Cancel operation requires a node-argv ref") 127 if operation == "cancel" and requested_by is not None and not isinstance(requested_by, str): 128 raise DmlRepoError("Cancel operation requested_by must be a string or null") 129 try: 130 valid_result_ref = isinstance(payload.get("result_ref"), str) and Ref(payload["result_ref"]).ns() == "dag" 131 except (TypeError, ValueError): 132 valid_result_ref = False 133 if operation == "cleanup" and not valid_result_ref: 134 raise DmlRepoError("Cleanup operation requires a non-null result_ref") 135 executor = cls() 136 if operation == "cancel": 137 result = executor.cancel( 138 cache_key=cache_key, 139 execution_id=execution_id, 140 runnable=runnable, 141 state=adapter_state, 142 remote=remote, 143 scratch_uri=scratch_uri, 144 cancel_requested_by=requested_by, 145 argv_ref=argv_ref, 146 ) 147 elif operation == "cleanup": 148 result = executor.cleanup( 149 cache_key=cache_key, 150 execution_id=execution_id, 151 runnable=runnable, 152 state=adapter_state, 153 remote=remote, 154 scratch_uri=scratch_uri, 155 result_ref=payload["result_ref"], 156 ) 157 elif adapter_state is None: 158 result = executor.start( 159 cache_key=cache_key, 160 execution_id=execution_id, 161 runnable=runnable, 162 remote=remote, 163 scratch_uri=scratch_uri, 164 ) 165 else: 166 result = executor.poll( 167 cache_key=cache_key, 168 execution_id=execution_id, 169 runnable=runnable, 170 state=adapter_state, 171 remote=remote, 172 scratch_uri=scratch_uri, 173 ) 174 if not isinstance(result, dict): 175 raise DmlRepoError("Executor response must be a JSON object") 176 if "state" in result and "adapter_state" in result: 177 raise DmlRepoError("Executor response cannot contain both state and adapter_state") 178 if "state" in result: 179 result["adapter_state"] = result.pop("state") 180 return validate_adapter_response( 181 result, 182 success_status="cancelled" if operation == "cancel" else "success", 183 )
Dispatch an explicit adapter operation to the executor.
71class BatchExecutor(LambdaExecutorBase): 72 name = "batch" 73 74 @staticmethod 75 def _string(name: str, value: Any) -> str: 76 if not isinstance(value, str) or not value: 77 raise DmlRepoError(f"batch executor {name} must be a non-empty string") 78 return value 79 80 @staticmethod 81 def _int(name: str, value: Any, *, default: int, min_value: int = 0) -> int: 82 if value is None: 83 return default 84 if not isinstance(value, int) or value < min_value: 85 raise DmlRepoError(f"batch executor {name} must be an int >= {min_value}") 86 return value 87 88 @classmethod 89 def _image_uri(cls, value: Any) -> Uri: 90 if not isinstance(value, Uri): 91 raise DmlRepoError("batch executor image must be a Uri") 92 return value 93 94 @classmethod 95 def resolve_runnable(cls, uri, kwargs, sub): 96 if sub is None: 97 raise DmlRepoError("batch executor requires sub runnable") 98 unknown = sorted(set(kwargs.keys()) - {"lambda_uri", "image", "cpu", "memory", "gpu"}) 99 if unknown: 100 raise DmlRepoError(f"Unknown batch executor kwargs: {', '.join(unknown)}") 101 return Runnable( 102 target=Uri(cls._string("lambda_uri", kwargs.get("lambda_uri"))), 103 adapter="dml-lambda-adapter", 104 kwargs={ 105 "image": cls._image_uri(kwargs.get("image")), 106 "cpu": cls._int("cpu", kwargs.get("cpu"), default=DEFAULT_VCPU, min_value=1), 107 "memory": cls._int("memory", kwargs.get("memory"), default=DEFAULT_MEMORY, min_value=1), 108 "gpu": cls._int("gpu", kwargs.get("gpu"), default=DEFAULT_GPU, min_value=0), 109 }, 110 sub=sub, 111 ) 112 113 @staticmethod 114 def _client(*, max_attempts: int = _BATCH_POLL_MAX_ATTEMPTS): 115 return _batch_client("batch", max_attempts=max_attempts) 116 117 @classmethod 118 def _resource_requirements(cls, kwargs: dict[str, Any]) -> tuple[list[dict[str, str]], str]: 119 cpu = cls._int("cpu", kwargs.get("cpu"), default=DEFAULT_VCPU, min_value=1) 120 memory = cls._int("memory", kwargs.get("memory"), default=DEFAULT_MEMORY, min_value=1) 121 gpu = cls._int("gpu", kwargs.get("gpu"), default=DEFAULT_GPU, min_value=0) 122 reqs = [ 123 {"type": "MEMORY", "value": str(memory)}, 124 {"type": "VCPU", "value": str(cpu)}, 125 ] 126 queue_env = "CPU_QUEUE" 127 if gpu > 0: 128 reqs.append({"type": "GPU", "value": str(gpu)}) 129 queue_env = "GPU_QUEUE" 130 return reqs, cls._string(queue_env, os.environ.get(queue_env)) 131 132 def start( 133 self, 134 *, 135 cache_key: str, 136 execution_id: str, 137 runnable: dict[str, Any], 138 remote: dict[str, str], 139 scratch_uri: str, 140 ) -> dict[str, Any]: 141 sub = runnable.get("sub") 142 if sub is None: 143 raise DmlRepoError("batch executor start requires runnable with sub runnable") 144 input_uri = _scratch_uri(scratch_uri, "input.json") 145 output_uri = _scratch_uri(scratch_uri, "output.json") 146 payload = json.dumps( 147 { 148 "operation": "invoke", 149 "runnable": sub, 150 "cache_key": cache_key, 151 "execution_id": execution_id, 152 "remote": remote, 153 "scratch_uri": scratch_uri, 154 "adapter_state": None, 155 } 156 ) 157 _write_scratch_json(input_uri, payload, raw=True) 158 client = self._client(max_attempts=_BATCH_START_MAX_ATTEMPTS) 159 kwargs = runnable.get("kwargs", {}) 160 reqs, job_queue = self._resource_requirements(kwargs) 161 image = self._image_uri(kwargs.get("image")) 162 job_name = f"dml-batch-{cache_key}" 163 job_def = client.register_job_definition( 164 jobDefinitionName=job_name, 165 type="container", 166 containerProperties={ 167 "image": image, 168 "command": [sub["adapter"], "--poll", "-i", input_uri, "-o", output_uri], 169 "environment": [], 170 "jobRoleArn": self._string("BATCH_TASK_ROLE_ARN", os.environ.get("BATCH_TASK_ROLE_ARN")), 171 "resourceRequirements": reqs, 172 }, 173 )["jobDefinitionArn"] 174 job_id = client.submit_job(jobName=job_name, jobQueue=job_queue, jobDefinition=job_def)["jobId"] 175 return { 176 "status": "retry", 177 "error": None, 178 "state": { 179 "job_id": job_id, 180 "job_definition": job_def, 181 }, 182 } 183 184 def poll( 185 self, 186 cache_key: str, 187 execution_id: str, 188 runnable: dict[str, Any], 189 state: dict[str, Any], 190 remote: dict[str, str], 191 scratch_uri: str, 192 ) -> dict[str, Any]: 193 del cache_key, execution_id, runnable, remote 194 job_id = state.get("job_id") 195 if not isinstance(job_id, str) or not job_id: 196 return { 197 "status": "failure", 198 "error": "batch poll: missing job_id in job state", 199 "state": state, 200 } 201 try: 202 jobs = self._client().describe_jobs(jobs=[job_id]).get("jobs", []) 203 except Exception as exc: 204 if self._is_throttling(exc): 205 result = {"status": "retry", "error": None, "state": state} 206 retry_after = self._retry_after(exc) 207 if retry_after is not None: 208 result["retry_after_ms"] = retry_after 209 return result 210 return {"status": "failure", "error": f"batch status check failed: {exc}", "state": state} 211 if not jobs: 212 return {"status": "retry", "error": None, "state": state} 213 job = jobs[0] 214 job_status = job["status"] 215 216 if job_status in PENDING_BATCH_STATUSES: 217 return {"status": "retry", "error": None, "state": state} 218 219 if job_status == "SUCCEEDED": 220 try: 221 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 222 if raw is None: 223 return { 224 "status": "failure", 225 "error": "batch poll: sub-adapter output not yet written to S3", 226 "state": state, 227 } 228 result = json.loads(raw) 229 except Exception as e: 230 return { 231 "status": "failure", 232 "error": f"batch poll: could not read sub-adapter result: {e}", 233 "state": state, 234 } 235 try: 236 result = validate_adapter_response(result) 237 except DmlRepoError as exc: 238 raise DmlRepoError(f"batch poll: invalid nested adapter output: {result}") from exc 239 nested_state = result.pop("adapter_state", None) 240 next_state = {**state, "nested_adapter_state": nested_state} if isinstance(nested_state, dict) else state 241 return {**result, "state": next_state} 242 243 # Failed 244 reason = None 245 if isinstance(job.get("statusReason"), str) and job["statusReason"]: 246 reason = job["statusReason"] 247 attempts = job.get("attempts") or [{}] 248 container = attempts[-1].get("container", {}) if attempts else {} 249 if isinstance(container, dict): 250 reason = container.get("reason") or container.get("exitCode") or reason 251 error = f"Batch job {job_id} failed" 252 if reason not in {None, ""}: 253 error = f"{error}: {reason}" 254 return {"status": "failure", "error": error, "state": state} 255 256 @staticmethod 257 def _is_throttling(exc: Exception) -> bool: 258 return getattr(exc, "response", {}).get("Error", {}).get("Code") in _THROTTLING_CODES 259 260 @staticmethod 261 def _retry_after(exc: Exception) -> int | None: 262 headers = getattr(exc, "response", {}).get("ResponseMetadata", {}).get("HTTPHeaders", {}) 263 value = headers.get("retry-after") 264 try: 265 return max(0, int(float(value) * 1000)) 266 except (TypeError, ValueError): 267 return None 268 269 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 270 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 271 state = state if isinstance(state, dict) else {} 272 job_id = state.get("job_id") 273 job_definition = state.get("job_definition") 274 client = self._client() 275 if isinstance(job_id, str) and job_id: 276 try: 277 jobs = client.describe_jobs(jobs=[job_id]).get("jobs", []) 278 except Exception as exc: 279 if self._is_throttling(exc): 280 result = {"status": "retry", "error": None, "state": state} 281 retry_after = self._retry_after(exc) 282 if retry_after is not None: 283 result["retry_after_ms"] = retry_after 284 return result 285 return {"status": "failure", "error": f"batch cleanup status check failed: {exc}", "state": state} 286 if jobs and jobs[0].get("status") in PENDING_BATCH_STATUSES: 287 return {"status": "retry", "error": None, "state": state} 288 if isinstance(job_definition, str) and job_definition: 289 try: 290 client.deregister_job_definition(jobDefinition=job_definition) 291 except Exception as exc: 292 if self._is_throttling(exc): 293 result = {"status": "retry", "error": None, "state": state} 294 retry_after = self._retry_after(exc) 295 if retry_after is not None: 296 result["retry_after_ms"] = retry_after 297 return result 298 return {"status": "failure", "error": f"batch cleanup failed: {exc}", "state": state} 299 return {"status": "success", "error": None, "state": state} 300 301 def cancel( 302 self, 303 cache_key: str, 304 execution_id: str, 305 runnable: dict[str, Any], 306 state: dict[str, Any], 307 remote: dict[str, str], 308 scratch_uri: str, 309 cancel_requested_by: str | None, 310 argv_ref: str | None = None, 311 ) -> dict[str, Any]: 312 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 313 client = self._client() 314 job_id = state.get("job_id") 315 job_definition = state.get("job_definition") 316 if isinstance(job_id, str) and job_id: 317 try: 318 client.cancel_job(jobId=job_id, reason="daggerml cancellation requested") 319 except Exception as cancel_exc: 320 try: 321 client.terminate_job(jobId=job_id, reason="daggerml cancellation requested") 322 except Exception as terminate_exc: 323 exc = terminate_exc if self._is_throttling(terminate_exc) else cancel_exc 324 if self._is_throttling(exc): 325 result = {"status": "retry", "error": None, "state": state} 326 retry_after = self._retry_after(exc) 327 if retry_after is not None: 328 result["retry_after_ms"] = retry_after 329 return result 330 return {"status": "failure", "error": f"batch cancellation failed: {terminate_exc}", "state": state} 331 if isinstance(job_definition, str) and job_definition: 332 try: 333 client.deregister_job_definition(jobDefinition=job_definition) 334 except Exception as exc: 335 return {"status": "failure", "error": f"batch cancellation failed: {exc}", "state": state} 336 return {"status": "cancelled", "error": None, "state": state}
Base class for all executors.
The runtime owns durable adapter state. Executors receive adapter_state=None
on first launch and persisted state on later status checks. Executors
return terminal or in-progress result dicts via stdout/return value:
{"status": "retry", "error": None, "state": {...}}
{"status": "success", "error": None, "state": None}
{"status": "failure", "error": "<msg>", "state": None}
94 @classmethod 95 def resolve_runnable(cls, uri, kwargs, sub): 96 if sub is None: 97 raise DmlRepoError("batch executor requires sub runnable") 98 unknown = sorted(set(kwargs.keys()) - {"lambda_uri", "image", "cpu", "memory", "gpu"}) 99 if unknown: 100 raise DmlRepoError(f"Unknown batch executor kwargs: {', '.join(unknown)}") 101 return Runnable( 102 target=Uri(cls._string("lambda_uri", kwargs.get("lambda_uri"))), 103 adapter="dml-lambda-adapter", 104 kwargs={ 105 "image": cls._image_uri(kwargs.get("image")), 106 "cpu": cls._int("cpu", kwargs.get("cpu"), default=DEFAULT_VCPU, min_value=1), 107 "memory": cls._int("memory", kwargs.get("memory"), default=DEFAULT_MEMORY, min_value=1), 108 "gpu": cls._int("gpu", kwargs.get("gpu"), default=DEFAULT_GPU, min_value=0), 109 }, 110 sub=sub, 111 )
132 def start( 133 self, 134 *, 135 cache_key: str, 136 execution_id: str, 137 runnable: dict[str, Any], 138 remote: dict[str, str], 139 scratch_uri: str, 140 ) -> dict[str, Any]: 141 sub = runnable.get("sub") 142 if sub is None: 143 raise DmlRepoError("batch executor start requires runnable with sub runnable") 144 input_uri = _scratch_uri(scratch_uri, "input.json") 145 output_uri = _scratch_uri(scratch_uri, "output.json") 146 payload = json.dumps( 147 { 148 "operation": "invoke", 149 "runnable": sub, 150 "cache_key": cache_key, 151 "execution_id": execution_id, 152 "remote": remote, 153 "scratch_uri": scratch_uri, 154 "adapter_state": None, 155 } 156 ) 157 _write_scratch_json(input_uri, payload, raw=True) 158 client = self._client(max_attempts=_BATCH_START_MAX_ATTEMPTS) 159 kwargs = runnable.get("kwargs", {}) 160 reqs, job_queue = self._resource_requirements(kwargs) 161 image = self._image_uri(kwargs.get("image")) 162 job_name = f"dml-batch-{cache_key}" 163 job_def = client.register_job_definition( 164 jobDefinitionName=job_name, 165 type="container", 166 containerProperties={ 167 "image": image, 168 "command": [sub["adapter"], "--poll", "-i", input_uri, "-o", output_uri], 169 "environment": [], 170 "jobRoleArn": self._string("BATCH_TASK_ROLE_ARN", os.environ.get("BATCH_TASK_ROLE_ARN")), 171 "resourceRequirements": reqs, 172 }, 173 )["jobDefinitionArn"] 174 job_id = client.submit_job(jobName=job_name, jobQueue=job_queue, jobDefinition=job_def)["jobId"] 175 return { 176 "status": "retry", 177 "error": None, 178 "state": { 179 "job_id": job_id, 180 "job_definition": job_def, 181 }, 182 }
Launch execution and return a result dict.
For synchronous executors this should return the terminal result
immediately. For async executors, return the durable resume state in the
initial retry result.
184 def poll( 185 self, 186 cache_key: str, 187 execution_id: str, 188 runnable: dict[str, Any], 189 state: dict[str, Any], 190 remote: dict[str, str], 191 scratch_uri: str, 192 ) -> dict[str, Any]: 193 del cache_key, execution_id, runnable, remote 194 job_id = state.get("job_id") 195 if not isinstance(job_id, str) or not job_id: 196 return { 197 "status": "failure", 198 "error": "batch poll: missing job_id in job state", 199 "state": state, 200 } 201 try: 202 jobs = self._client().describe_jobs(jobs=[job_id]).get("jobs", []) 203 except Exception as exc: 204 if self._is_throttling(exc): 205 result = {"status": "retry", "error": None, "state": state} 206 retry_after = self._retry_after(exc) 207 if retry_after is not None: 208 result["retry_after_ms"] = retry_after 209 return result 210 return {"status": "failure", "error": f"batch status check failed: {exc}", "state": state} 211 if not jobs: 212 return {"status": "retry", "error": None, "state": state} 213 job = jobs[0] 214 job_status = job["status"] 215 216 if job_status in PENDING_BATCH_STATUSES: 217 return {"status": "retry", "error": None, "state": state} 218 219 if job_status == "SUCCEEDED": 220 try: 221 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 222 if raw is None: 223 return { 224 "status": "failure", 225 "error": "batch poll: sub-adapter output not yet written to S3", 226 "state": state, 227 } 228 result = json.loads(raw) 229 except Exception as e: 230 return { 231 "status": "failure", 232 "error": f"batch poll: could not read sub-adapter result: {e}", 233 "state": state, 234 } 235 try: 236 result = validate_adapter_response(result) 237 except DmlRepoError as exc: 238 raise DmlRepoError(f"batch poll: invalid nested adapter output: {result}") from exc 239 nested_state = result.pop("adapter_state", None) 240 next_state = {**state, "nested_adapter_state": nested_state} if isinstance(nested_state, dict) else state 241 return {**result, "state": next_state} 242 243 # Failed 244 reason = None 245 if isinstance(job.get("statusReason"), str) and job["statusReason"]: 246 reason = job["statusReason"] 247 attempts = job.get("attempts") or [{}] 248 container = attempts[-1].get("container", {}) if attempts else {} 249 if isinstance(container, dict): 250 reason = container.get("reason") or container.get("exitCode") or reason 251 error = f"Batch job {job_id} failed" 252 if reason not in {None, ""}: 253 error = f"{error}: {reason}" 254 return {"status": "failure", "error": error, "state": state}
Check an in-flight job and return a result dict.
state is the immutable launch-time state returned by start().
Return a terminal result when done, or {"status": "retry",
"error": None, "state": ...} while still running.
Later returned state may be ignored by the runtime.
269 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 270 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 271 state = state if isinstance(state, dict) else {} 272 job_id = state.get("job_id") 273 job_definition = state.get("job_definition") 274 client = self._client() 275 if isinstance(job_id, str) and job_id: 276 try: 277 jobs = client.describe_jobs(jobs=[job_id]).get("jobs", []) 278 except Exception as exc: 279 if self._is_throttling(exc): 280 result = {"status": "retry", "error": None, "state": state} 281 retry_after = self._retry_after(exc) 282 if retry_after is not None: 283 result["retry_after_ms"] = retry_after 284 return result 285 return {"status": "failure", "error": f"batch cleanup status check failed: {exc}", "state": state} 286 if jobs and jobs[0].get("status") in PENDING_BATCH_STATUSES: 287 return {"status": "retry", "error": None, "state": state} 288 if isinstance(job_definition, str) and job_definition: 289 try: 290 client.deregister_job_definition(jobDefinition=job_definition) 291 except Exception as exc: 292 if self._is_throttling(exc): 293 result = {"status": "retry", "error": None, "state": state} 294 retry_after = self._retry_after(exc) 295 if retry_after is not None: 296 result["retry_after_ms"] = retry_after 297 return result 298 return {"status": "failure", "error": f"batch cleanup failed: {exc}", "state": state} 299 return {"status": "success", "error": None, "state": state}
Idempotently prune resources after a result was published.
301 def cancel( 302 self, 303 cache_key: str, 304 execution_id: str, 305 runnable: dict[str, Any], 306 state: dict[str, Any], 307 remote: dict[str, str], 308 scratch_uri: str, 309 cancel_requested_by: str | None, 310 argv_ref: str | None = None, 311 ) -> dict[str, Any]: 312 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 313 client = self._client() 314 job_id = state.get("job_id") 315 job_definition = state.get("job_definition") 316 if isinstance(job_id, str) and job_id: 317 try: 318 client.cancel_job(jobId=job_id, reason="daggerml cancellation requested") 319 except Exception as cancel_exc: 320 try: 321 client.terminate_job(jobId=job_id, reason="daggerml cancellation requested") 322 except Exception as terminate_exc: 323 exc = terminate_exc if self._is_throttling(terminate_exc) else cancel_exc 324 if self._is_throttling(exc): 325 result = {"status": "retry", "error": None, "state": state} 326 retry_after = self._retry_after(exc) 327 if retry_after is not None: 328 result["retry_after_ms"] = retry_after 329 return result 330 return {"status": "failure", "error": f"batch cancellation failed: {terminate_exc}", "state": state} 331 if isinstance(job_definition, str) and job_definition: 332 try: 333 client.deregister_job_definition(jobDefinition=job_definition) 334 except Exception as exc: 335 return {"status": "failure", "error": f"batch cancellation failed: {exc}", "state": state} 336 return {"status": "cancelled", "error": None, "state": state}
56class DockerExecutor(ExecutorBase): 57 name = "docker" 58 adapter = "local" 59 60 @classmethod 61 def resolve_runnable(cls, uri, kwargs, sub): 62 if sub is None: 63 raise DmlRepoError("docker executor requires sub runnable") 64 image = kwargs.get("image") 65 if image is None: 66 raise DmlRepoError("docker executor requires image") 67 unknown = sorted(set(kwargs.keys()) - {"image", "flags"}) 68 if unknown: 69 raise DmlRepoError(f"Unknown docker executor kwargs: {', '.join(unknown)}") 70 return Runnable( 71 target=Uri("docker"), 72 kwargs={"image": image, "flags": kwargs.get("flags", [])}, 73 sub=sub, 74 adapter="dml-local-adapter", 75 ) 76 77 @staticmethod 78 def _run_docker(*args: str, check: bool = True, docker_bin: str | None = None) -> str: 79 docker_bin = docker_bin or shutil.which("docker") 80 if docker_bin is None: 81 raise DmlRepoError("docker executable not found in PATH") 82 proc = subprocess.run([docker_bin, *args], check=False, capture_output=True, text=True) 83 if proc.returncode == 0: 84 return proc.stdout.strip() or proc.stderr.strip() 85 if check: 86 command = f"{docker_bin} {' '.join(args)}" 87 raise DmlRepoError( 88 f"docker command failed ({proc.returncode}): {command}\nSTDOUT:\n{proc.stdout}\nSTDERR:\n{proc.stderr}" 89 ) 90 return proc.stdout.strip() or proc.stderr.strip() 91 92 @staticmethod 93 def _image_input(runnable: dict[str, Any]) -> str: 94 image = runnable.get("kwargs", {}).get("image") 95 if isinstance(image, dict): 96 image = image.get("uri") 97 if isinstance(image, Uri): 98 return image.uri 99 if isinstance(image, str) and image: 100 return image 101 raise DmlRepoError("docker executor image must resolve to a non-empty Uri or string") 102 103 @staticmethod 104 def _image_tag_from_tar(tar_path: Path) -> str: 105 with tarfile.open(tar_path, mode="r") as tf: 106 member = tf.extractfile("manifest.json") 107 if member is None: 108 raise DmlRepoError("docker image tar missing manifest.json") 109 manifest = json.loads(member.read()) 110 repo_tags = manifest[0].get("RepoTags") if manifest else None 111 if not isinstance(repo_tags, list) or not repo_tags or not isinstance(repo_tags[0], str) or not repo_tags[0]: 112 raise DmlRepoError("docker image tar missing RepoTags") 113 return cast(str, repo_tags[0]) 114 115 @staticmethod 116 def _isolate_image_tar(tar_path: Path, isolated_path: Path) -> str: 117 DockerExecutor._image_tag_from_tar(tar_path) 118 isolated_tag = f"daggerml-execution:{uuid4().hex}" 119 with tarfile.open(tar_path, mode="r:*") as source, tarfile.open(isolated_path, mode="w") as target: 120 for member in source: 121 if member.name == "repositories": 122 continue # Legacy Docker metadata would also restore the original tag. 123 content = source.extractfile(member) if member.isfile() else None 124 if member.name == "manifest.json": 125 assert content is not None 126 manifest = json.load(content) 127 if not isinstance(manifest, list) or len(manifest) != 1: 128 raise DmlRepoError("docker image tar must contain exactly one image") 129 manifest[0]["RepoTags"] = [isolated_tag] 130 data = json.dumps(manifest).encode("utf-8") 131 member = copy.copy(member) 132 member.size = len(data) 133 content = io.BytesIO(data) 134 target.addfile(member, content) 135 return isolated_tag 136 137 @staticmethod 138 def _prepare_image(runnable: dict[str, Any], workdir: Path, remote: dict[str, Any]) -> tuple[str, str | None]: 139 image = DockerExecutor._image_input(runnable) 140 if not is_s3_uri(image): 141 return image, None 142 tar_path = workdir / "image.tar" 143 store = S3Store.from_remote_root(cast(str, remote["root"])) 144 tar_path.write_bytes(store.get(image)) 145 isolated_path = workdir / "isolated.tar" 146 image_ref = DockerExecutor._isolate_image_tar(tar_path, isolated_path) 147 DockerExecutor._run_docker("load", "-i", str(isolated_path)) 148 return image_ref, image_ref 149 150 def start( 151 self, 152 *, 153 cache_key: str, 154 execution_id: str, 155 runnable: dict[str, Any], 156 remote: dict[str, str], 157 scratch_uri: str, 158 ) -> dict[str, Any]: 159 sub = runnable.get("sub") 160 if sub is None: 161 raise DmlRepoError("docker executor requires sub runnable") 162 input_uri = _scratch_uri(scratch_uri, "input.json") 163 output_uri = _scratch_uri(scratch_uri, "output.json") 164 165 workdir = Path(tempfile.mkdtemp(prefix=f"dml-docker-{execution_id}-")) 166 try: 167 image_ref, cleanup_image = self._prepare_image(runnable, workdir, remote) 168 finally: 169 shutil.rmtree(workdir, ignore_errors=True) 170 171 payload = json.dumps( 172 { 173 "operation": "invoke", 174 "runnable": sub, 175 "cache_key": cache_key, 176 "execution_id": execution_id, 177 "remote": remote, 178 "scratch_uri": scratch_uri, 179 "adapter_state": None, 180 } 181 ) 182 _write_scratch_json(input_uri, payload, raw=True) 183 184 container_id = self._run_docker( 185 "run", 186 "-d", 187 *cast(list[str], runnable.get("kwargs", {}).get("flags", [])), 188 "-e", 189 f"DML_REMOTE_ROOT={remote['root']}", 190 image_ref, 191 sub["adapter"], 192 "--poll", 193 "-i", 194 input_uri, 195 "-o", 196 output_uri, 197 ) 198 199 return { 200 "status": "retry", 201 "error": None, 202 "state": { 203 "owner": _current_owner(), 204 "container_id": container_id, 205 "cleanup_image": cleanup_image, 206 }, 207 } 208 209 def poll( 210 self, 211 cache_key: str, 212 execution_id: str, 213 runnable: dict[str, Any], 214 state: dict[str, Any], 215 remote: dict[str, str], 216 scratch_uri: str, 217 ) -> dict[str, Any]: 218 del cache_key, execution_id, runnable, remote 219 if state["owner"] != _current_owner(): 220 return {"status": "retry", "error": None, "state": state} 221 container_id = state.get("container_id") 222 223 if not isinstance(container_id, str) or not container_id: 224 return { 225 "status": "failure", 226 "error": "docker poll: missing container_id in job state", 227 "state": state, 228 } 229 230 docker_bin = shutil.which("docker") 231 if docker_bin is None: 232 return { 233 "status": "failure", 234 "error": "docker poll: docker executable not found", 235 "state": state, 236 } 237 238 proc = subprocess.run( 239 [docker_bin, "inspect", "--format", "{{.State.Status}}", container_id], 240 check=False, 241 capture_output=True, 242 text=True, 243 ) 244 if proc.returncode != 0: 245 # An inspect error cannot establish that the container has exited. 246 return {"status": "retry", "error": None, "state": state} 247 else: 248 container_status = proc.stdout.strip() 249 250 if container_status in ("created", "running", "paused", "restarting"): 251 return {"status": "retry", "error": None, "state": state} 252 253 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 254 if raw is not None: 255 try: 256 result = validate_adapter_response(json.loads(raw)) 257 result.pop("adapter_state", None) 258 return {**result, "state": state} 259 except Exception as e: 260 raise DmlRepoError(f"docker poll: invalid nested adapter output: {e}") from e 261 262 return { 263 "status": "failure", 264 "error": f"docker container {container_id} exited without output", 265 "state": state, 266 } 267 268 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 269 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 270 if isinstance(state, dict) and state["owner"] != _current_owner(): 271 return {"status": "retry", "error": None, "state": state} 272 state = state if isinstance(state, dict) else {} 273 container_id = state.get("container_id") 274 docker_bin = shutil.which("docker") 275 if not isinstance(container_id, str) or not container_id or docker_bin is None: 276 return {"status": "success", "error": None, "state": state} 277 status = self._run_docker( 278 "inspect", "--format", "{{.State.Status}}", container_id, check=False, docker_bin=docker_bin 279 ) 280 if status in {"created", "running", "paused", "restarting"}: 281 return {"status": "retry", "error": None, "state": state} 282 _cleanup_docker(container_id, state.get("cleanup_image"), docker_bin) 283 return {"status": "success", "error": None, "state": state} 284 285 def cancel( 286 self, 287 cache_key: str, 288 execution_id: str, 289 runnable: dict[str, Any], 290 state: dict[str, Any], 291 remote: dict[str, str], 292 scratch_uri: str, 293 cancel_requested_by: str | None, 294 argv_ref: str | None = None, 295 ) -> dict[str, Any]: 296 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 297 if isinstance(state, dict) and state["owner"] != _current_owner(): 298 return {"status": "retry", "error": None, "state": state} 299 state = state if isinstance(state, dict) else {} 300 docker_bin = shutil.which("docker") 301 container_id = state.get("container_id") 302 if docker_bin is None: 303 if isinstance(container_id, str) and container_id: 304 return {"status": "failure", "error": "docker executable not found in PATH", "state": state} 305 return {"status": "cancelled", "error": None, "state": state} 306 if isinstance(container_id, str) and container_id: 307 _cleanup_docker(container_id, state.get("cleanup_image"), docker_bin) 308 return {"status": "cancelled", "error": None, "state": state}
Base class for all executors.
The runtime owns durable adapter state. Executors receive adapter_state=None
on first launch and persisted state on later status checks. Executors
return terminal or in-progress result dicts via stdout/return value:
{"status": "retry", "error": None, "state": {...}}
{"status": "success", "error": None, "state": None}
{"status": "failure", "error": "<msg>", "state": None}
60 @classmethod 61 def resolve_runnable(cls, uri, kwargs, sub): 62 if sub is None: 63 raise DmlRepoError("docker executor requires sub runnable") 64 image = kwargs.get("image") 65 if image is None: 66 raise DmlRepoError("docker executor requires image") 67 unknown = sorted(set(kwargs.keys()) - {"image", "flags"}) 68 if unknown: 69 raise DmlRepoError(f"Unknown docker executor kwargs: {', '.join(unknown)}") 70 return Runnable( 71 target=Uri("docker"), 72 kwargs={"image": image, "flags": kwargs.get("flags", [])}, 73 sub=sub, 74 adapter="dml-local-adapter", 75 )
150 def start( 151 self, 152 *, 153 cache_key: str, 154 execution_id: str, 155 runnable: dict[str, Any], 156 remote: dict[str, str], 157 scratch_uri: str, 158 ) -> dict[str, Any]: 159 sub = runnable.get("sub") 160 if sub is None: 161 raise DmlRepoError("docker executor requires sub runnable") 162 input_uri = _scratch_uri(scratch_uri, "input.json") 163 output_uri = _scratch_uri(scratch_uri, "output.json") 164 165 workdir = Path(tempfile.mkdtemp(prefix=f"dml-docker-{execution_id}-")) 166 try: 167 image_ref, cleanup_image = self._prepare_image(runnable, workdir, remote) 168 finally: 169 shutil.rmtree(workdir, ignore_errors=True) 170 171 payload = json.dumps( 172 { 173 "operation": "invoke", 174 "runnable": sub, 175 "cache_key": cache_key, 176 "execution_id": execution_id, 177 "remote": remote, 178 "scratch_uri": scratch_uri, 179 "adapter_state": None, 180 } 181 ) 182 _write_scratch_json(input_uri, payload, raw=True) 183 184 container_id = self._run_docker( 185 "run", 186 "-d", 187 *cast(list[str], runnable.get("kwargs", {}).get("flags", [])), 188 "-e", 189 f"DML_REMOTE_ROOT={remote['root']}", 190 image_ref, 191 sub["adapter"], 192 "--poll", 193 "-i", 194 input_uri, 195 "-o", 196 output_uri, 197 ) 198 199 return { 200 "status": "retry", 201 "error": None, 202 "state": { 203 "owner": _current_owner(), 204 "container_id": container_id, 205 "cleanup_image": cleanup_image, 206 }, 207 }
Launch execution and return a result dict.
For synchronous executors this should return the terminal result
immediately. For async executors, return the durable resume state in the
initial retry result.
209 def poll( 210 self, 211 cache_key: str, 212 execution_id: str, 213 runnable: dict[str, Any], 214 state: dict[str, Any], 215 remote: dict[str, str], 216 scratch_uri: str, 217 ) -> dict[str, Any]: 218 del cache_key, execution_id, runnable, remote 219 if state["owner"] != _current_owner(): 220 return {"status": "retry", "error": None, "state": state} 221 container_id = state.get("container_id") 222 223 if not isinstance(container_id, str) or not container_id: 224 return { 225 "status": "failure", 226 "error": "docker poll: missing container_id in job state", 227 "state": state, 228 } 229 230 docker_bin = shutil.which("docker") 231 if docker_bin is None: 232 return { 233 "status": "failure", 234 "error": "docker poll: docker executable not found", 235 "state": state, 236 } 237 238 proc = subprocess.run( 239 [docker_bin, "inspect", "--format", "{{.State.Status}}", container_id], 240 check=False, 241 capture_output=True, 242 text=True, 243 ) 244 if proc.returncode != 0: 245 # An inspect error cannot establish that the container has exited. 246 return {"status": "retry", "error": None, "state": state} 247 else: 248 container_status = proc.stdout.strip() 249 250 if container_status in ("created", "running", "paused", "restarting"): 251 return {"status": "retry", "error": None, "state": state} 252 253 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 254 if raw is not None: 255 try: 256 result = validate_adapter_response(json.loads(raw)) 257 result.pop("adapter_state", None) 258 return {**result, "state": state} 259 except Exception as e: 260 raise DmlRepoError(f"docker poll: invalid nested adapter output: {e}") from e 261 262 return { 263 "status": "failure", 264 "error": f"docker container {container_id} exited without output", 265 "state": state, 266 }
Check an in-flight job and return a result dict.
state is the immutable launch-time state returned by start().
Return a terminal result when done, or {"status": "retry",
"error": None, "state": ...} while still running.
Later returned state may be ignored by the runtime.
268 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 269 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 270 if isinstance(state, dict) and state["owner"] != _current_owner(): 271 return {"status": "retry", "error": None, "state": state} 272 state = state if isinstance(state, dict) else {} 273 container_id = state.get("container_id") 274 docker_bin = shutil.which("docker") 275 if not isinstance(container_id, str) or not container_id or docker_bin is None: 276 return {"status": "success", "error": None, "state": state} 277 status = self._run_docker( 278 "inspect", "--format", "{{.State.Status}}", container_id, check=False, docker_bin=docker_bin 279 ) 280 if status in {"created", "running", "paused", "restarting"}: 281 return {"status": "retry", "error": None, "state": state} 282 _cleanup_docker(container_id, state.get("cleanup_image"), docker_bin) 283 return {"status": "success", "error": None, "state": state}
Idempotently prune resources after a result was published.
285 def cancel( 286 self, 287 cache_key: str, 288 execution_id: str, 289 runnable: dict[str, Any], 290 state: dict[str, Any], 291 remote: dict[str, str], 292 scratch_uri: str, 293 cancel_requested_by: str | None, 294 argv_ref: str | None = None, 295 ) -> dict[str, Any]: 296 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 297 if isinstance(state, dict) and state["owner"] != _current_owner(): 298 return {"status": "retry", "error": None, "state": state} 299 state = state if isinstance(state, dict) else {} 300 docker_bin = shutil.which("docker") 301 container_id = state.get("container_id") 302 if docker_bin is None: 303 if isinstance(container_id, str) and container_id: 304 return {"status": "failure", "error": "docker executable not found in PATH", "state": state} 305 return {"status": "cancelled", "error": None, "state": state} 306 if isinstance(container_id, str) and container_id: 307 _cleanup_docker(container_id, state.get("cleanup_image"), docker_bin) 308 return {"status": "cancelled", "error": None, "state": state}
33class ScriptExecutor(ExecutorBase): 34 name = "script" 35 adapter = "local" 36 37 ############################# resolve runnable ############################# 38 ############################################################################ 39 40 @classmethod 41 def resolve_runnable(cls, uri, kwargs, sub): 42 if sub is not None: 43 raise DmlRepoError("script executor does not accept sub runnable") 44 resolved_kw, script = cls._script_kwargs(dict(kwargs)) 45 resolved_kw["script_uri"] = S3Store().put(data=script.encode("utf-8"), suffix=".py") 46 return Runnable(target=Uri("script"), kwargs=resolved_kw, sub=sub, adapter="dml-local-adapter") 47 48 @classmethod 49 def _script_kwargs(cls, kwargs: dict) -> tuple[dict, str]: 50 allowed = {"fn", "fn_name", "script", "prepop", "extra_objs", "post_lines", "tags"} 51 unknown = sorted(set(kwargs.keys()) - allowed) 52 if unknown: 53 bad = ", ".join(unknown) 54 raise DmlRepoError(f"Unknown script executor kwargs: {bad}") 55 fn = kwargs.get("fn") 56 prepop = kwargs.get("prepop", {}) 57 extra_objs = list(kwargs.get("extra_objs", [])) 58 post_lines = list(kwargs.get("post_lines", [])) 59 tags = kwargs.get("tags", []) 60 if not isinstance(tags, list) or not all(isinstance(tag, str) for tag in tags): 61 raise DmlRepoError("script tags must be a list of strings") 62 script = kwargs.get("script") 63 fn_name = kwargs.get("fn_name") 64 if (script is None) != (fn_name is None): 65 raise DmlRepoError("script and fn_name must be captured together") 66 if script is None: 67 if not callable(fn): 68 raise DmlRepoError("script fn must be callable") 69 params = list(inspect.signature(fn).parameters.values()) 70 if not params: 71 raise DmlRepoError("script fn must include at least one parameter") 72 script = cls._render_script(fn, extra_objs=extra_objs, post_lines=post_lines) 73 fn_name = fn.__name__ 74 elif not isinstance(script, str) or not isinstance(fn_name, str): 75 raise DmlRepoError("captured script and fn_name must be strings") 76 resolved = {"prepop": prepop, "fn_name": fn_name} 77 if tags: 78 resolved["tags"] = sorted(set(tags)) 79 return resolved, script 80 81 @classmethod 82 def _render_script(cls, fn, extra_objs: list, post_lines: list[str]) -> str: 83 chunks: list[str] = [] 84 for obj in [*extra_objs, fn]: 85 try: 86 raw = dedent(inspect.getsource(inspect.unwrap(obj))).strip() 87 chunks.append(cls._strip_funkify_decorators(raw)) 88 except (OSError, TypeError) as e: 89 raise DmlRepoError(f"Failed to serialize object source: {e}") from e 90 if post_lines: 91 chunks.extend(post_lines) 92 script = "\n\n".join(chunks) 93 try: 94 mod = ast.parse(script) 95 except SyntaxError as e: 96 raise DmlRepoError(f"Generated script is not valid Python: {e}") from e 97 if not any(isinstance(n, ast.FunctionDef) and n.name == fn.__name__ for n in mod.body): 98 raise DmlRepoError(f"Function '{fn.__name__}' is not globally defined in generated script") 99 return script 100 101 @staticmethod 102 def _strip_funkify_decorators(source: str) -> str: 103 module = ast.parse(source) 104 for node in module.body: 105 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): 106 node.decorator_list = [] 107 return ast.unparse(module).strip() 108 109 ############################# start/poll/cancel ############################ 110 111 def start( 112 self, 113 cache_key: str, 114 execution_id: str, 115 runnable: dict[str, Any], 116 remote: dict[str, str], 117 scratch_uri: str, 118 ) -> AdapterInvokeResponse: 119 del scratch_uri 120 tags = runnable.get("kwargs", {}).get("tags") 121 workdir = Path(tempfile.mkdtemp(prefix=f"dml-script-{execution_id}-")) 122 payload_path = workdir / "supervisor-input.json" 123 result_path = workdir / "result.json" 124 stdout_path = workdir / "stdout.log" 125 stderr_path = workdir / "stderr.log" 126 cmd = [ 127 sys.executable, 128 "-m", 129 "daggerml.contrib.executors.script", 130 "--execution-id", 131 execution_id, 132 "--cache-key", 133 cache_key, 134 "--remote-root", 135 remote["root"], 136 ] 137 if tags: 138 cmd.extend(["--tags", json.dumps(tags, separators=(",", ":"))]) 139 payload = { 140 "version": 0, 141 "cache_key": cache_key, 142 "execution_id": execution_id, 143 "cmd": cmd, 144 "remote": remote, 145 "env": {}, 146 } 147 payload_path.write_text(json.dumps(payload, separators=(",", ":"), sort_keys=True)) 148 with stdout_path.open("w") as stdout_f, stderr_path.open("w") as stderr_f: 149 proc = subprocess.Popen( 150 [ 151 sys.executable, 152 "-m", 153 "daggerml.contrib.supervisor", 154 "-i", 155 str(payload_path), 156 "-o", 157 str(result_path), 158 ], 159 stdout=stdout_f, 160 stderr=stderr_f, 161 start_new_session=True, 162 close_fds=True, 163 env={**os.environ, "PYTHONUNBUFFERED": "1", **payload["env"]}, 164 ) 165 launch_state = { 166 "owner": _current_owner(), 167 "pid": proc.pid, 168 "workdir": str(workdir), 169 "result_path": str(result_path), 170 "stdout_path": str(stdout_path), 171 "stderr_path": str(stderr_path), 172 } 173 return {"status": "retry", "error": None, "state": launch_state} 174 175 def poll( 176 self, 177 cache_key: str, 178 execution_id: str, 179 runnable: dict[str, Any], 180 state: dict[str, Any], 181 remote: dict[str, str], 182 scratch_uri: str, 183 ) -> AdapterInvokeResponse: 184 del cache_key, execution_id, runnable, remote, scratch_uri 185 if state["owner"] != _current_owner(): 186 return {"status": "retry", "error": None, "state": state} 187 result_path = Path(state.get("result_path", "")) 188 pid = state.get("pid") 189 # Polls may run either in the launching adapter process or in a later 190 # process. Reap children when we can; otherwise fall back to a direct 191 # PID probe for cross-process polling. 192 if isinstance(pid, int): 193 try: 194 done_pid, _ = os.waitpid(pid, os.WNOHANG) 195 if done_pid == 0: 196 return {"status": "retry", "error": None, "state": state} 197 except ChildProcessError: 198 try: 199 os.kill(pid, 0) 200 return {"status": "retry", "error": None, "state": state} 201 except ProcessLookupError: 202 pass 203 except PermissionError: 204 return {"status": "retry", "error": None, "state": state} 205 # Process exited — read result 206 if result_path.exists(): 207 try: 208 parsed = json.loads(result_path.read_text()) 209 if parsed.get("status") in {"succeeded", "failed"}: 210 return { 211 "status": "success" if parsed["status"] == "succeeded" else "failure", 212 "error": parsed.get("error"), 213 "state": state, 214 } 215 except Exception as e: 216 return { 217 "status": "failure", 218 "error": f"Could not read supervisor result: {e}", 219 "state": state, 220 } 221 return { 222 "status": "failure", 223 "error": "Script supervisor exited without result", 224 "state": state, 225 } 226 227 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 228 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 229 if isinstance(state, dict) and state["owner"] != _current_owner(): 230 return {"status": "retry", "error": None, "state": state} 231 state = state if isinstance(state, dict) else {} 232 pid = state.get("pid") 233 if isinstance(pid, int): 234 try: 235 done_pid, _ = os.waitpid(pid, os.WNOHANG) 236 active = done_pid == 0 237 except ChildProcessError: 238 try: 239 os.kill(pid, 0) 240 active = True 241 except ProcessLookupError: 242 active = False 243 except PermissionError as exc: 244 return {"status": "failure", "error": f"script cleanup failed: {exc}", "state": state} 245 if active: 246 # The DML result is already published, so the supervisor is no longer needed. 247 # Stop its process group before removing scratch so it cannot keep writing there. 248 try: 249 os.killpg(pid, signal.SIGTERM) 250 except ProcessLookupError: 251 pass 252 except PermissionError as exc: 253 return {"status": "failure", "error": f"script cleanup failed: {exc}", "state": state} 254 _cleanup_workdir(state) 255 workdir = state.get("workdir") 256 if isinstance(workdir, str) and os.path.exists(workdir): 257 return {"status": "failure", "error": f"script cleanup failed to remove {workdir}", "state": state} 258 return {"status": "success", "error": None, "state": state} 259 260 def cancel( 261 self, 262 cache_key: str, 263 execution_id: str, 264 runnable: dict[str, Any], 265 state: dict[str, Any], 266 remote: dict[str, str], 267 scratch_uri: str, 268 cancel_requested_by: str | None, 269 argv_ref: str | None = None, 270 ) -> AdapterCancelResponse: 271 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 272 if not isinstance(state, dict): 273 return {"status": "cancelled", "error": None, "state": {}} 274 if state["owner"] != _current_owner(): 275 return {"status": "retry", "error": None, "state": state} 276 pid = state.get("pid") 277 if isinstance(pid, int): 278 try: 279 os.killpg(pid, signal.SIGTERM) 280 except ProcessLookupError: 281 pass 282 except PermissionError as exc: 283 return {"status": "failure", "error": f"script cancellation failed: {exc}", "state": state} 284 _cleanup_workdir(state) 285 workdir = state.get("workdir") 286 if isinstance(workdir, str) and os.path.exists(workdir): 287 return {"status": "failure", "error": f"script cancellation failed to remove {workdir}", "state": state} 288 return {"status": "cancelled", "error": None, "state": state}
Base class for all executors.
The runtime owns durable adapter state. Executors receive adapter_state=None
on first launch and persisted state on later status checks. Executors
return terminal or in-progress result dicts via stdout/return value:
{"status": "retry", "error": None, "state": {...}}
{"status": "success", "error": None, "state": None}
{"status": "failure", "error": "<msg>", "state": None}
40 @classmethod 41 def resolve_runnable(cls, uri, kwargs, sub): 42 if sub is not None: 43 raise DmlRepoError("script executor does not accept sub runnable") 44 resolved_kw, script = cls._script_kwargs(dict(kwargs)) 45 resolved_kw["script_uri"] = S3Store().put(data=script.encode("utf-8"), suffix=".py") 46 return Runnable(target=Uri("script"), kwargs=resolved_kw, sub=sub, adapter="dml-local-adapter")
111 def start( 112 self, 113 cache_key: str, 114 execution_id: str, 115 runnable: dict[str, Any], 116 remote: dict[str, str], 117 scratch_uri: str, 118 ) -> AdapterInvokeResponse: 119 del scratch_uri 120 tags = runnable.get("kwargs", {}).get("tags") 121 workdir = Path(tempfile.mkdtemp(prefix=f"dml-script-{execution_id}-")) 122 payload_path = workdir / "supervisor-input.json" 123 result_path = workdir / "result.json" 124 stdout_path = workdir / "stdout.log" 125 stderr_path = workdir / "stderr.log" 126 cmd = [ 127 sys.executable, 128 "-m", 129 "daggerml.contrib.executors.script", 130 "--execution-id", 131 execution_id, 132 "--cache-key", 133 cache_key, 134 "--remote-root", 135 remote["root"], 136 ] 137 if tags: 138 cmd.extend(["--tags", json.dumps(tags, separators=(",", ":"))]) 139 payload = { 140 "version": 0, 141 "cache_key": cache_key, 142 "execution_id": execution_id, 143 "cmd": cmd, 144 "remote": remote, 145 "env": {}, 146 } 147 payload_path.write_text(json.dumps(payload, separators=(",", ":"), sort_keys=True)) 148 with stdout_path.open("w") as stdout_f, stderr_path.open("w") as stderr_f: 149 proc = subprocess.Popen( 150 [ 151 sys.executable, 152 "-m", 153 "daggerml.contrib.supervisor", 154 "-i", 155 str(payload_path), 156 "-o", 157 str(result_path), 158 ], 159 stdout=stdout_f, 160 stderr=stderr_f, 161 start_new_session=True, 162 close_fds=True, 163 env={**os.environ, "PYTHONUNBUFFERED": "1", **payload["env"]}, 164 ) 165 launch_state = { 166 "owner": _current_owner(), 167 "pid": proc.pid, 168 "workdir": str(workdir), 169 "result_path": str(result_path), 170 "stdout_path": str(stdout_path), 171 "stderr_path": str(stderr_path), 172 } 173 return {"status": "retry", "error": None, "state": launch_state}
Launch execution and return a result dict.
For synchronous executors this should return the terminal result
immediately. For async executors, return the durable resume state in the
initial retry result.
175 def poll( 176 self, 177 cache_key: str, 178 execution_id: str, 179 runnable: dict[str, Any], 180 state: dict[str, Any], 181 remote: dict[str, str], 182 scratch_uri: str, 183 ) -> AdapterInvokeResponse: 184 del cache_key, execution_id, runnable, remote, scratch_uri 185 if state["owner"] != _current_owner(): 186 return {"status": "retry", "error": None, "state": state} 187 result_path = Path(state.get("result_path", "")) 188 pid = state.get("pid") 189 # Polls may run either in the launching adapter process or in a later 190 # process. Reap children when we can; otherwise fall back to a direct 191 # PID probe for cross-process polling. 192 if isinstance(pid, int): 193 try: 194 done_pid, _ = os.waitpid(pid, os.WNOHANG) 195 if done_pid == 0: 196 return {"status": "retry", "error": None, "state": state} 197 except ChildProcessError: 198 try: 199 os.kill(pid, 0) 200 return {"status": "retry", "error": None, "state": state} 201 except ProcessLookupError: 202 pass 203 except PermissionError: 204 return {"status": "retry", "error": None, "state": state} 205 # Process exited — read result 206 if result_path.exists(): 207 try: 208 parsed = json.loads(result_path.read_text()) 209 if parsed.get("status") in {"succeeded", "failed"}: 210 return { 211 "status": "success" if parsed["status"] == "succeeded" else "failure", 212 "error": parsed.get("error"), 213 "state": state, 214 } 215 except Exception as e: 216 return { 217 "status": "failure", 218 "error": f"Could not read supervisor result: {e}", 219 "state": state, 220 } 221 return { 222 "status": "failure", 223 "error": "Script supervisor exited without result", 224 "state": state, 225 }
Check an in-flight job and return a result dict.
state is the immutable launch-time state returned by start().
Return a terminal result when done, or {"status": "retry",
"error": None, "state": ...} while still running.
Later returned state may be ignored by the runtime.
227 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 228 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 229 if isinstance(state, dict) and state["owner"] != _current_owner(): 230 return {"status": "retry", "error": None, "state": state} 231 state = state if isinstance(state, dict) else {} 232 pid = state.get("pid") 233 if isinstance(pid, int): 234 try: 235 done_pid, _ = os.waitpid(pid, os.WNOHANG) 236 active = done_pid == 0 237 except ChildProcessError: 238 try: 239 os.kill(pid, 0) 240 active = True 241 except ProcessLookupError: 242 active = False 243 except PermissionError as exc: 244 return {"status": "failure", "error": f"script cleanup failed: {exc}", "state": state} 245 if active: 246 # The DML result is already published, so the supervisor is no longer needed. 247 # Stop its process group before removing scratch so it cannot keep writing there. 248 try: 249 os.killpg(pid, signal.SIGTERM) 250 except ProcessLookupError: 251 pass 252 except PermissionError as exc: 253 return {"status": "failure", "error": f"script cleanup failed: {exc}", "state": state} 254 _cleanup_workdir(state) 255 workdir = state.get("workdir") 256 if isinstance(workdir, str) and os.path.exists(workdir): 257 return {"status": "failure", "error": f"script cleanup failed to remove {workdir}", "state": state} 258 return {"status": "success", "error": None, "state": state}
Idempotently prune resources after a result was published.
260 def cancel( 261 self, 262 cache_key: str, 263 execution_id: str, 264 runnable: dict[str, Any], 265 state: dict[str, Any], 266 remote: dict[str, str], 267 scratch_uri: str, 268 cancel_requested_by: str | None, 269 argv_ref: str | None = None, 270 ) -> AdapterCancelResponse: 271 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 272 if not isinstance(state, dict): 273 return {"status": "cancelled", "error": None, "state": {}} 274 if state["owner"] != _current_owner(): 275 return {"status": "retry", "error": None, "state": state} 276 pid = state.get("pid") 277 if isinstance(pid, int): 278 try: 279 os.killpg(pid, signal.SIGTERM) 280 except ProcessLookupError: 281 pass 282 except PermissionError as exc: 283 return {"status": "failure", "error": f"script cancellation failed: {exc}", "state": state} 284 _cleanup_workdir(state) 285 workdir = state.get("workdir") 286 if isinstance(workdir, str) and os.path.exists(workdir): 287 return {"status": "failure", "error": f"script cancellation failed to remove {workdir}", "state": state} 288 return {"status": "cancelled", "error": None, "state": state}
25class SshExecutor(ExecutorBase): 26 name = "ssh" 27 adapter = "local" 28 29 def start( 30 self, 31 cache_key: str, 32 execution_id: str, 33 runnable: dict[str, Any], 34 remote: dict[str, str], 35 scratch_uri: str, 36 ) -> dict[str, Any]: 37 return self._send_nested( 38 cache_key=cache_key, 39 execution_id=execution_id, 40 runnable=runnable, 41 remote=remote, 42 scratch_uri=scratch_uri, 43 operation="invoke", 44 adapter_state=None, 45 cancel_requested_by=None, 46 ) 47 48 def poll( 49 self, 50 cache_key: str, 51 execution_id: str, 52 runnable: dict[str, Any], 53 state: dict[str, Any], 54 remote: dict[str, str], 55 scratch_uri: str, 56 ) -> dict[str, Any]: 57 return self._send_nested( 58 cache_key=cache_key, 59 execution_id=execution_id, 60 runnable=runnable, 61 remote=remote, 62 scratch_uri=scratch_uri, 63 operation="invoke", 64 adapter_state=state, 65 cancel_requested_by=None, 66 ) 67 68 def cancel( 69 self, 70 cache_key: str, 71 execution_id: str, 72 runnable: dict[str, Any], 73 state: dict[str, Any], 74 remote: dict[str, str], 75 scratch_uri: str, 76 cancel_requested_by: str | None, 77 argv_ref: str | None = None, 78 ) -> dict[str, Any]: 79 return self._send_nested( 80 cache_key=cache_key, 81 execution_id=execution_id, 82 runnable=runnable, 83 remote=remote, 84 scratch_uri=scratch_uri, 85 operation="cancel", 86 adapter_state=state, 87 cancel_requested_by=cancel_requested_by, 88 argv_ref=argv_ref, 89 ) 90 91 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 92 return self._send_nested( 93 cache_key=cache_key, 94 execution_id=execution_id, 95 runnable=runnable, 96 remote=remote, 97 scratch_uri=scratch_uri, 98 operation="cleanup", 99 adapter_state=state, 100 result_ref=result_ref, 101 cancel_requested_by=None, 102 ) 103 104 @classmethod 105 def _send_nested( 106 cls, 107 *, 108 cache_key: str, 109 execution_id: str, 110 runnable: dict[str, Any], 111 remote: dict[str, str], 112 scratch_uri: str, 113 operation: str, 114 adapter_state: dict[str, Any] | None, 115 cancel_requested_by: str | None, 116 argv_ref: str | None = None, 117 result_ref: str | None = None, 118 ) -> dict[str, Any]: 119 sub = runnable.get("sub") 120 if sub is None: 121 raise DmlRepoError("ssh executor requires sub runnable") 122 kw = cls._validate_kw(cast(dict, runnable.get("kwargs", {}))) 123 cmd = [ 124 "ssh", 125 *kw["flags"], 126 kw["host"], 127 cls._remote_command(env_files=kw["env_files"], adapter=sub["adapter"]), 128 ] 129 payload = { 130 "operation": operation, 131 "runnable": sub, 132 "cache_key": cache_key, 133 "execution_id": execution_id, 134 "remote": remote, 135 "scratch_uri": scratch_uri, 136 "adapter_state": adapter_state, 137 } 138 if operation == "cancel": 139 payload["requested_by"] = cancel_requested_by 140 payload["argv_ref"] = argv_ref 141 elif operation == "cleanup": 142 payload["result_ref"] = result_ref 143 payload = json.dumps(payload) 144 logger.debug( 145 "ssh executor launch host=%s flags=%s env_files=%s adapter=%s cache_key=%s execution_id=%s has_state=%s", 146 kw["host"], 147 kw["flags"], 148 kw["env_files"], 149 sub["adapter"], 150 cache_key, 151 execution_id, 152 adapter_state is not None, 153 ) 154 proc = subprocess.run(cmd, input=payload, capture_output=True, check=False, text=True) 155 stdout = proc.stdout.strip() 156 stderr = proc.stderr.strip() 157 logger.debug( 158 "ssh executor command returncode=%s execution_id=%s stdout=%r stderr=%r", 159 proc.returncode, 160 execution_id, 161 stdout, 162 stderr, 163 ) 164 if proc.returncode != 0: 165 error = f"SSH command failed ({proc.returncode})" 166 if stderr: 167 error = f"{error}: {stderr}" 168 elif stdout: 169 error = f"{error}: {stdout}" 170 logger.debug("ssh executor transport failed execution_id=%s error=%s", execution_id, error) 171 return {"status": "failure", "error": error, "adapter_state": adapter_state or {}} 172 try: 173 result = json.loads(stdout) 174 except json.JSONDecodeError as e: 175 logger.debug( 176 "ssh executor invalid json execution_id=%s error=%s stdout=%r", 177 execution_id, 178 e, 179 stdout, 180 ) 181 return { 182 "status": "failure", 183 "error": f"SSH nested adapter returned invalid JSON: {e}", 184 "adapter_state": adapter_state or {}, 185 } 186 try: 187 result = validate_adapter_response( 188 result, 189 success_status="cancelled" if operation == "cancel" else "success", 190 ) 191 except DmlRepoError as exc: 192 logger.debug("ssh executor unexpected result execution_id=%s result=%r", execution_id, result) 193 raise DmlRepoError(f"SSH nested adapter returned unexpected result: {exc}") from exc 194 logger.debug( 195 "ssh executor result execution_id=%s status=%s error=%r", 196 execution_id, 197 result.get("status"), 198 result.get("error"), 199 ) 200 return result 201 202 @staticmethod 203 def _validate_kw(kw: dict) -> SshExecKwargs: 204 if not isinstance(kw, dict): 205 raise DmlRepoError("ssh executor kwargs must be a dict") 206 if set(kw.keys()) > {"env_files", "flags", "host"}: 207 raise DmlRepoError("ssh executor kwargs only supports keys: env_files, flags, host") 208 host = cast(str, kw.get("host")) 209 if not (is_node_like(host) or (isinstance(host, str) and host)): 210 raise DmlRepoError("ssh executor requires non-empty host") 211 kw["flags"] = flags = cast(list[str], kw.get("flags") or []) 212 if not _is_node_string_list(flags): 213 raise DmlRepoError("ssh executor flags must be a list of non-empty strings") 214 kw["env_files"] = env_files = cast(list[str], kw.get("env_files") or []) 215 if not _is_node_string_list(env_files): 216 raise DmlRepoError("ssh executor env_files must be a list of non-empty strings") 217 return SshExecKwargs(host=host, flags=flags, env_files=env_files) 218 219 @classmethod 220 def resolve_runnable(cls, uri, kwargs, sub): 221 if sub is None: 222 raise DmlRepoError("ssh executor requires sub runnable") 223 unknown = sorted(set(kwargs.keys()) - {"env_files", "flags", "host"}) 224 if unknown: 225 raise DmlRepoError(f"Unknown ssh executor kwargs: {', '.join(unknown)}") 226 return Runnable( 227 target=Uri("ssh"), 228 kwargs=dict(cls._validate_kw(kwargs)), 229 sub=sub, 230 adapter="dml-local-adapter", 231 ) 232 233 @staticmethod 234 def _remote_command(*, env_files: list[str], adapter: str) -> str: 235 parts = ["set -e"] 236 parts.extend(f". {shlex.quote(path)}" for path in env_files) 237 parts.append(f"exec {shlex.quote(adapter)} -i - -o -") 238 return "; ".join(parts)
Base class for all executors.
The runtime owns durable adapter state. Executors receive adapter_state=None
on first launch and persisted state on later status checks. Executors
return terminal or in-progress result dicts via stdout/return value:
{"status": "retry", "error": None, "state": {...}}
{"status": "success", "error": None, "state": None}
{"status": "failure", "error": "<msg>", "state": None}
29 def start( 30 self, 31 cache_key: str, 32 execution_id: str, 33 runnable: dict[str, Any], 34 remote: dict[str, str], 35 scratch_uri: str, 36 ) -> dict[str, Any]: 37 return self._send_nested( 38 cache_key=cache_key, 39 execution_id=execution_id, 40 runnable=runnable, 41 remote=remote, 42 scratch_uri=scratch_uri, 43 operation="invoke", 44 adapter_state=None, 45 cancel_requested_by=None, 46 )
Launch execution and return a result dict.
For synchronous executors this should return the terminal result
immediately. For async executors, return the durable resume state in the
initial retry result.
48 def poll( 49 self, 50 cache_key: str, 51 execution_id: str, 52 runnable: dict[str, Any], 53 state: dict[str, Any], 54 remote: dict[str, str], 55 scratch_uri: str, 56 ) -> dict[str, Any]: 57 return self._send_nested( 58 cache_key=cache_key, 59 execution_id=execution_id, 60 runnable=runnable, 61 remote=remote, 62 scratch_uri=scratch_uri, 63 operation="invoke", 64 adapter_state=state, 65 cancel_requested_by=None, 66 )
Check an in-flight job and return a result dict.
state is the immutable launch-time state returned by start().
Return a terminal result when done, or {"status": "retry",
"error": None, "state": ...} while still running.
Later returned state may be ignored by the runtime.
68 def cancel( 69 self, 70 cache_key: str, 71 execution_id: str, 72 runnable: dict[str, Any], 73 state: dict[str, Any], 74 remote: dict[str, str], 75 scratch_uri: str, 76 cancel_requested_by: str | None, 77 argv_ref: str | None = None, 78 ) -> dict[str, Any]: 79 return self._send_nested( 80 cache_key=cache_key, 81 execution_id=execution_id, 82 runnable=runnable, 83 remote=remote, 84 scratch_uri=scratch_uri, 85 operation="cancel", 86 adapter_state=state, 87 cancel_requested_by=cancel_requested_by, 88 argv_ref=argv_ref, 89 )
91 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 92 return self._send_nested( 93 cache_key=cache_key, 94 execution_id=execution_id, 95 runnable=runnable, 96 remote=remote, 97 scratch_uri=scratch_uri, 98 operation="cleanup", 99 adapter_state=state, 100 result_ref=result_ref, 101 cancel_requested_by=None, 102 )
Idempotently prune resources after a result was published.
219 @classmethod 220 def resolve_runnable(cls, uri, kwargs, sub): 221 if sub is None: 222 raise DmlRepoError("ssh executor requires sub runnable") 223 unknown = sorted(set(kwargs.keys()) - {"env_files", "flags", "host"}) 224 if unknown: 225 raise DmlRepoError(f"Unknown ssh executor kwargs: {', '.join(unknown)}") 226 return Runnable( 227 target=Uri("ssh"), 228 kwargs=dict(cls._validate_kw(kwargs)), 229 sub=sub, 230 adapter="dml-local-adapter", 231 )