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.
72class BatchExecutor(LambdaExecutorBase): 73 name = "batch" 74 75 @staticmethod 76 def _string(name: str, value: Any) -> str: 77 if not isinstance(value, str) or not value: 78 raise DmlRepoError(f"batch executor {name} must be a non-empty string") 79 return value 80 81 @staticmethod 82 def _int(name: str, value: Any, *, default: int, min_value: int = 0) -> int: 83 if value is None: 84 return default 85 if isinstance(value, bool) or not isinstance(value, int) or value < min_value: 86 raise DmlRepoError(f"batch executor {name} must be an int >= {min_value}") 87 return value 88 89 @classmethod 90 def _image_uri(cls, value: Any) -> Uri: 91 if not isinstance(value, Uri): 92 raise DmlRepoError("batch executor image must be a Uri") 93 cls._string("image URI", value.uri) 94 return value 95 96 @classmethod 97 def resolve_runnable(cls, uri, kwargs, sub): 98 if sub is None: 99 raise DmlRepoError("batch executor requires sub runnable") 100 unknown = sorted(set(kwargs.keys()) - {"lambda_uri", "image", "cpu", "memory", "gpu"}) 101 if unknown: 102 raise DmlRepoError(f"Unknown batch executor kwargs: {', '.join(unknown)}") 103 return Runnable( 104 target=Uri(cls._string("lambda_uri", kwargs.get("lambda_uri"))), 105 adapter="dml-lambda-adapter", 106 kwargs={ 107 "image": cls._image_uri(kwargs.get("image")), 108 "cpu": cls._int("cpu", kwargs.get("cpu"), default=DEFAULT_VCPU, min_value=1), 109 "memory": cls._int("memory", kwargs.get("memory"), default=DEFAULT_MEMORY, min_value=1), 110 "gpu": cls._int("gpu", kwargs.get("gpu"), default=DEFAULT_GPU, min_value=0), 111 }, 112 sub=sub, 113 ) 114 115 @staticmethod 116 def _client(*, max_attempts: int = _BATCH_POLL_MAX_ATTEMPTS): 117 return _batch_client("batch", max_attempts=max_attempts) 118 119 @classmethod 120 def _resource_requirements(cls, kwargs: dict[str, Any]) -> tuple[list[dict[str, str]], str]: 121 cpu = cls._int("cpu", kwargs.get("cpu"), default=DEFAULT_VCPU, min_value=1) 122 memory = cls._int("memory", kwargs.get("memory"), default=DEFAULT_MEMORY, min_value=1) 123 gpu = cls._int("gpu", kwargs.get("gpu"), default=DEFAULT_GPU, min_value=0) 124 reqs = [ 125 {"type": "MEMORY", "value": str(memory)}, 126 {"type": "VCPU", "value": str(cpu)}, 127 ] 128 queue_env = "CPU_QUEUE" 129 if gpu > 0: 130 reqs.append({"type": "GPU", "value": str(gpu)}) 131 queue_env = "GPU_QUEUE" 132 return reqs, cls._string(queue_env, os.environ.get(queue_env)) 133 134 def start( 135 self, 136 *, 137 cache_key: str, 138 execution_id: str, 139 runnable: dict[str, Any], 140 remote: dict[str, str], 141 scratch_uri: str, 142 ) -> dict[str, Any]: 143 return self._launch(cache_key, execution_id, runnable, {}, remote, scratch_uri) 144 145 def _launch(self, cache_key, execution_id, runnable, state, remote, scratch_uri) -> dict[str, Any]: 146 state = dict(state) 147 try: 148 return self._start(cache_key, execution_id, runnable, state, remote, scratch_uri) 149 except Exception as exc: 150 if not self._is_throttling(exc): 151 raise 152 result = {"status": "retry", "error": None, "state": state} 153 retry_after = self._retry_after(exc) 154 if retry_after is not None: 155 result["retry_after_ms"] = retry_after 156 return result 157 158 def _start(self, cache_key, execution_id, runnable, state, remote, scratch_uri) -> dict[str, Any]: 159 sub = runnable.get("sub") 160 if sub is None: 161 raise DmlRepoError("batch executor start requires runnable with sub runnable") 162 kwargs = runnable.get("kwargs", {}) 163 reqs, job_queue = self._resource_requirements(kwargs) 164 image = kwargs.get("image") 165 # Runtime requests carry JSON URI objects, not authoring-time dataclasses. 166 if isinstance(image, dict) and set(image) == {"uri"}: 167 image = Uri(image["uri"]) 168 image = self._image_uri(image).uri 169 adapter = self._string("sub adapter", sub.get("adapter")) 170 task_role = self._string("BATCH_TASK_ROLE_ARN", os.environ.get("BATCH_TASK_ROLE_ARN")) 171 input_uri = _scratch_uri(scratch_uri, "input.json") 172 output_uri = _scratch_uri(scratch_uri, "output.json") 173 payload = json.dumps( 174 { 175 "operation": "invoke", 176 "runnable": sub, 177 "cache_key": cache_key, 178 "execution_id": execution_id, 179 "remote": remote, 180 "scratch_uri": scratch_uri, 181 "adapter_state": None, 182 } 183 ) 184 _write_scratch_json(input_uri, payload, raw=True) 185 client = self._client(max_attempts=_BATCH_START_MAX_ATTEMPTS) 186 region = self._string("AWS region", client.meta.region_name) 187 job_name = f"dml-batch-{hashlib.sha256(execution_id.encode()).hexdigest()}" 188 job_def = state.get("job_definition") or client.register_job_definition( 189 jobDefinitionName=job_name, 190 type="container", 191 containerProperties={ 192 "image": image, 193 "command": [adapter, "--poll", "-i", input_uri, "-o", output_uri], 194 "environment": [ 195 {"name": "AWS_EC2_METADATA_DISABLED", "value": "true"}, 196 {"name": "AWS_REGION", "value": region}, 197 {"name": "AWS_DEFAULT_REGION", "value": region}, 198 ], 199 "privileged": False, 200 "jobRoleArn": task_role, 201 "resourceRequirements": reqs, 202 }, 203 )["jobDefinitionArn"] 204 state["job_definition"] = job_def 205 try: 206 job_id = client.submit_job(jobName=job_name, jobQueue=job_queue, jobDefinition=job_def)["jobId"] 207 except Exception as exc: 208 if self._is_throttling(exc): 209 # A rejected submission can resume using the registered definition. 210 raise 211 try: 212 client.deregister_job_definition(jobDefinition=job_def) 213 except Exception as cleanup_exc: 214 return { 215 "status": "failure", 216 "error": f"batch submission failed: {exc}; rollback failed for {job_def}: {cleanup_exc}", 217 "state": {"job_definition": job_def}, 218 } 219 raise 220 return { 221 "status": "retry", 222 "error": None, 223 "state": { 224 "job_id": job_id, 225 "job_definition": job_def, 226 }, 227 } 228 229 def poll( 230 self, 231 cache_key: str, 232 execution_id: str, 233 runnable: dict[str, Any], 234 state: dict[str, Any], 235 remote: dict[str, str], 236 scratch_uri: str, 237 ) -> dict[str, Any]: 238 if not state or (state.get("job_definition") and not state.get("job_id")): 239 return self._launch(cache_key, execution_id, runnable, state, remote, scratch_uri) 240 del cache_key, execution_id, runnable, remote 241 job_id = state.get("job_id") 242 if not isinstance(job_id, str) or not job_id: 243 return { 244 "status": "failure", 245 "error": "batch poll: missing job_id in job state", 246 "state": state, 247 } 248 try: 249 jobs = self._client().describe_jobs(jobs=[job_id]).get("jobs", []) 250 except Exception as exc: 251 if self._is_throttling(exc): 252 result = {"status": "retry", "error": None, "state": state} 253 retry_after = self._retry_after(exc) 254 if retry_after is not None: 255 result["retry_after_ms"] = retry_after 256 return result 257 return {"status": "failure", "error": f"batch status check failed: {exc}", "state": state} 258 if not jobs: 259 return {"status": "retry", "error": None, "state": state} 260 job = jobs[0] 261 job_status = job["status"] 262 263 if job_status in PENDING_BATCH_STATUSES: 264 return {"status": "retry", "error": None, "state": state} 265 266 if job_status == "SUCCEEDED": 267 try: 268 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 269 if raw is None: 270 return { 271 "status": "failure", 272 "error": "batch poll: sub-adapter output not yet written to S3", 273 "state": state, 274 } 275 result = json.loads(raw) 276 except Exception as e: 277 return { 278 "status": "failure", 279 "error": f"batch poll: could not read sub-adapter result: {e}", 280 "state": state, 281 } 282 try: 283 result = validate_adapter_response(result) 284 except DmlRepoError as exc: 285 raise DmlRepoError(f"batch poll: invalid nested adapter output: {result}") from exc 286 nested_state = result.pop("adapter_state", None) 287 next_state = {**state, "nested_adapter_state": nested_state} if isinstance(nested_state, dict) else state 288 return {**result, "state": next_state} 289 290 # Failed 291 reason = None 292 if isinstance(job.get("statusReason"), str) and job["statusReason"]: 293 reason = job["statusReason"] 294 attempts = job.get("attempts") or [{}] 295 container = attempts[-1].get("container", {}) if attempts else {} 296 if isinstance(container, dict): 297 reason = container.get("reason") or container.get("exitCode") or reason 298 error = f"Batch job {job_id} failed" 299 if reason not in {None, ""}: 300 error = f"{error}: {reason}" 301 return {"status": "failure", "error": error, "state": state} 302 303 @staticmethod 304 def _is_throttling(exc: Exception) -> bool: 305 return getattr(exc, "response", {}).get("Error", {}).get("Code") in _THROTTLING_CODES 306 307 @staticmethod 308 def _retry_after(exc: Exception) -> int | None: 309 headers = getattr(exc, "response", {}).get("ResponseMetadata", {}).get("HTTPHeaders", {}) 310 value = headers.get("retry-after") 311 try: 312 return max(0, int(float(value) * 1000)) 313 except (TypeError, ValueError): 314 return None 315 316 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 317 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 318 state = state if isinstance(state, dict) else {} 319 job_id = state.get("job_id") 320 job_definition = state.get("job_definition") 321 client = self._client() 322 if isinstance(job_id, str) and job_id: 323 try: 324 jobs = client.describe_jobs(jobs=[job_id]).get("jobs", []) 325 except Exception as exc: 326 if self._is_throttling(exc): 327 result = {"status": "retry", "error": None, "state": state} 328 retry_after = self._retry_after(exc) 329 if retry_after is not None: 330 result["retry_after_ms"] = retry_after 331 return result 332 return {"status": "failure", "error": f"batch cleanup status check failed: {exc}", "state": state} 333 if jobs and jobs[0].get("status") in PENDING_BATCH_STATUSES: 334 return {"status": "retry", "error": None, "state": state} 335 if isinstance(job_definition, str) and job_definition: 336 try: 337 client.deregister_job_definition(jobDefinition=job_definition) 338 except Exception as exc: 339 if self._is_throttling(exc): 340 result = {"status": "retry", "error": None, "state": state} 341 retry_after = self._retry_after(exc) 342 if retry_after is not None: 343 result["retry_after_ms"] = retry_after 344 return result 345 return {"status": "failure", "error": f"batch cleanup failed: {exc}", "state": state} 346 return {"status": "success", "error": None, "state": state} 347 348 def cancel( 349 self, 350 cache_key: str, 351 execution_id: str, 352 runnable: dict[str, Any], 353 state: dict[str, Any], 354 remote: dict[str, str], 355 scratch_uri: str, 356 cancel_requested_by: str | None, 357 argv_ref: str | None = None, 358 ) -> dict[str, Any]: 359 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 360 state = state if isinstance(state, dict) else {} 361 client = self._client() 362 job_id = state.get("job_id") 363 job_definition = state.get("job_definition") 364 if isinstance(job_id, str) and job_id: 365 try: 366 # TerminateJob also cancels queued jobs; CancelJob is a no-op once started. 367 client.terminate_job(jobId=job_id, reason="daggerml cancellation requested") 368 jobs = client.describe_jobs(jobs=[job_id]).get("jobs", []) 369 if not jobs or jobs[0].get("status") not in {"SUCCEEDED", "FAILED"}: 370 return {"status": "retry", "error": None, "state": state} 371 except Exception as exc: 372 if self._is_throttling(exc): 373 result = {"status": "retry", "error": None, "state": state} 374 retry_after = self._retry_after(exc) 375 if retry_after is not None: 376 result["retry_after_ms"] = retry_after 377 return result 378 return {"status": "failure", "error": f"batch cancellation failed: {exc}", "state": state} 379 if isinstance(job_definition, str) and job_definition: 380 try: 381 client.deregister_job_definition(jobDefinition=job_definition) 382 except Exception as exc: 383 if self._is_throttling(exc): 384 result = {"status": "retry", "error": None, "state": state} 385 retry_after = self._retry_after(exc) 386 if retry_after is not None: 387 result["retry_after_ms"] = retry_after 388 return result 389 return {"status": "failure", "error": f"batch cancellation failed: {exc}", "state": state} 390 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}
96 @classmethod 97 def resolve_runnable(cls, uri, kwargs, sub): 98 if sub is None: 99 raise DmlRepoError("batch executor requires sub runnable") 100 unknown = sorted(set(kwargs.keys()) - {"lambda_uri", "image", "cpu", "memory", "gpu"}) 101 if unknown: 102 raise DmlRepoError(f"Unknown batch executor kwargs: {', '.join(unknown)}") 103 return Runnable( 104 target=Uri(cls._string("lambda_uri", kwargs.get("lambda_uri"))), 105 adapter="dml-lambda-adapter", 106 kwargs={ 107 "image": cls._image_uri(kwargs.get("image")), 108 "cpu": cls._int("cpu", kwargs.get("cpu"), default=DEFAULT_VCPU, min_value=1), 109 "memory": cls._int("memory", kwargs.get("memory"), default=DEFAULT_MEMORY, min_value=1), 110 "gpu": cls._int("gpu", kwargs.get("gpu"), default=DEFAULT_GPU, min_value=0), 111 }, 112 sub=sub, 113 )
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.
229 def poll( 230 self, 231 cache_key: str, 232 execution_id: str, 233 runnable: dict[str, Any], 234 state: dict[str, Any], 235 remote: dict[str, str], 236 scratch_uri: str, 237 ) -> dict[str, Any]: 238 if not state or (state.get("job_definition") and not state.get("job_id")): 239 return self._launch(cache_key, execution_id, runnable, state, remote, scratch_uri) 240 del cache_key, execution_id, runnable, remote 241 job_id = state.get("job_id") 242 if not isinstance(job_id, str) or not job_id: 243 return { 244 "status": "failure", 245 "error": "batch poll: missing job_id in job state", 246 "state": state, 247 } 248 try: 249 jobs = self._client().describe_jobs(jobs=[job_id]).get("jobs", []) 250 except Exception as exc: 251 if self._is_throttling(exc): 252 result = {"status": "retry", "error": None, "state": state} 253 retry_after = self._retry_after(exc) 254 if retry_after is not None: 255 result["retry_after_ms"] = retry_after 256 return result 257 return {"status": "failure", "error": f"batch status check failed: {exc}", "state": state} 258 if not jobs: 259 return {"status": "retry", "error": None, "state": state} 260 job = jobs[0] 261 job_status = job["status"] 262 263 if job_status in PENDING_BATCH_STATUSES: 264 return {"status": "retry", "error": None, "state": state} 265 266 if job_status == "SUCCEEDED": 267 try: 268 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 269 if raw is None: 270 return { 271 "status": "failure", 272 "error": "batch poll: sub-adapter output not yet written to S3", 273 "state": state, 274 } 275 result = json.loads(raw) 276 except Exception as e: 277 return { 278 "status": "failure", 279 "error": f"batch poll: could not read sub-adapter result: {e}", 280 "state": state, 281 } 282 try: 283 result = validate_adapter_response(result) 284 except DmlRepoError as exc: 285 raise DmlRepoError(f"batch poll: invalid nested adapter output: {result}") from exc 286 nested_state = result.pop("adapter_state", None) 287 next_state = {**state, "nested_adapter_state": nested_state} if isinstance(nested_state, dict) else state 288 return {**result, "state": next_state} 289 290 # Failed 291 reason = None 292 if isinstance(job.get("statusReason"), str) and job["statusReason"]: 293 reason = job["statusReason"] 294 attempts = job.get("attempts") or [{}] 295 container = attempts[-1].get("container", {}) if attempts else {} 296 if isinstance(container, dict): 297 reason = container.get("reason") or container.get("exitCode") or reason 298 error = f"Batch job {job_id} failed" 299 if reason not in {None, ""}: 300 error = f"{error}: {reason}" 301 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.
316 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 317 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 318 state = state if isinstance(state, dict) else {} 319 job_id = state.get("job_id") 320 job_definition = state.get("job_definition") 321 client = self._client() 322 if isinstance(job_id, str) and job_id: 323 try: 324 jobs = client.describe_jobs(jobs=[job_id]).get("jobs", []) 325 except Exception as exc: 326 if self._is_throttling(exc): 327 result = {"status": "retry", "error": None, "state": state} 328 retry_after = self._retry_after(exc) 329 if retry_after is not None: 330 result["retry_after_ms"] = retry_after 331 return result 332 return {"status": "failure", "error": f"batch cleanup status check failed: {exc}", "state": state} 333 if jobs and jobs[0].get("status") in PENDING_BATCH_STATUSES: 334 return {"status": "retry", "error": None, "state": state} 335 if isinstance(job_definition, str) and job_definition: 336 try: 337 client.deregister_job_definition(jobDefinition=job_definition) 338 except Exception as exc: 339 if self._is_throttling(exc): 340 result = {"status": "retry", "error": None, "state": state} 341 retry_after = self._retry_after(exc) 342 if retry_after is not None: 343 result["retry_after_ms"] = retry_after 344 return result 345 return {"status": "failure", "error": f"batch cleanup failed: {exc}", "state": state} 346 return {"status": "success", "error": None, "state": state}
Idempotently prune resources after a result was published.
348 def cancel( 349 self, 350 cache_key: str, 351 execution_id: str, 352 runnable: dict[str, Any], 353 state: dict[str, Any], 354 remote: dict[str, str], 355 scratch_uri: str, 356 cancel_requested_by: str | None, 357 argv_ref: str | None = None, 358 ) -> dict[str, Any]: 359 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 360 state = state if isinstance(state, dict) else {} 361 client = self._client() 362 job_id = state.get("job_id") 363 job_definition = state.get("job_definition") 364 if isinstance(job_id, str) and job_id: 365 try: 366 # TerminateJob also cancels queued jobs; CancelJob is a no-op once started. 367 client.terminate_job(jobId=job_id, reason="daggerml cancellation requested") 368 jobs = client.describe_jobs(jobs=[job_id]).get("jobs", []) 369 if not jobs or jobs[0].get("status") not in {"SUCCEEDED", "FAILED"}: 370 return {"status": "retry", "error": None, "state": state} 371 except Exception as exc: 372 if self._is_throttling(exc): 373 result = {"status": "retry", "error": None, "state": state} 374 retry_after = self._retry_after(exc) 375 if retry_after is not None: 376 result["retry_after_ms"] = retry_after 377 return result 378 return {"status": "failure", "error": f"batch cancellation failed: {exc}", "state": state} 379 if isinstance(job_definition, str) and job_definition: 380 try: 381 client.deregister_job_definition(jobDefinition=job_definition) 382 except Exception as exc: 383 if self._is_throttling(exc): 384 result = {"status": "retry", "error": None, "state": state} 385 retry_after = self._retry_after(exc) 386 if retry_after is not None: 387 result["retry_after_ms"] = retry_after 388 return result 389 return {"status": "failure", "error": f"batch cancellation failed: {exc}", "state": state} 390 return {"status": "cancelled", "error": None, "state": state}
84class DockerExecutor(ExecutorBase): 85 name = "docker" 86 adapter = "local" 87 88 @classmethod 89 def resolve_runnable(cls, uri, kwargs, sub): 90 if sub is None: 91 raise DmlRepoError("docker executor requires sub runnable") 92 image = kwargs.get("image") 93 if image is None: 94 raise DmlRepoError("docker executor requires image") 95 unknown = sorted(set(kwargs.keys()) - {"image", "flags"}) 96 if unknown: 97 raise DmlRepoError(f"Unknown docker executor kwargs: {', '.join(unknown)}") 98 return Runnable( 99 target=Uri("docker"), 100 kwargs={"image": image, "flags": kwargs.get("flags", [])}, 101 sub=sub, 102 adapter="dml-local-adapter", 103 ) 104 105 @staticmethod 106 def _run_docker(*args: str, check: bool = True, docker_bin: str | None = None) -> str: 107 docker_bin = docker_bin or shutil.which("docker") 108 if docker_bin is None: 109 raise DmlRepoError("docker executable not found in PATH") 110 proc = subprocess.run([docker_bin, *args], check=False, capture_output=True, text=True) 111 if proc.returncode == 0: 112 return proc.stdout.strip() or proc.stderr.strip() 113 if check: 114 command = f"{docker_bin} {' '.join(args)}" 115 raise DmlRepoError( 116 f"docker command failed ({proc.returncode}): {command}\nSTDOUT:\n{proc.stdout}\nSTDERR:\n{proc.stderr}" 117 ) 118 return proc.stdout.strip() or proc.stderr.strip() 119 120 @staticmethod 121 def _image_input(runnable: dict[str, Any]) -> str: 122 image = runnable.get("kwargs", {}).get("image") 123 if isinstance(image, dict): 124 image = image.get("uri") 125 if isinstance(image, Uri): 126 return image.uri 127 if isinstance(image, str) and image: 128 return image 129 raise DmlRepoError("docker executor image must resolve to a non-empty Uri or string") 130 131 @staticmethod 132 def _stage_image(runnable: dict[str, Any], workdir: Path, remote: dict[str, Any]) -> str | Path: 133 """Download archives before acquiring the daemon lifecycle lock.""" 134 image = DockerExecutor._image_input(runnable) 135 if not is_s3_uri(image): 136 return image 137 tar_path = workdir / "image.tar" 138 store = S3Store.from_remote_root(cast(str, remote["root"])) 139 tar_path.write_bytes(store.get(image)) 140 return tar_path 141 142 @staticmethod 143 def _prepare_image(image: str | Path) -> str: 144 """Resolve Docker's loaded image ID; caller holds the lock through launch.""" 145 if isinstance(image, str): 146 return image 147 output = DockerExecutor._run_docker("load", "-i", str(image)) 148 references = set() 149 for line in output.splitlines(): 150 for prefix in ("Loaded image: ", "Loaded image ID: "): 151 if line.strip().startswith(prefix): 152 reference = line.strip().removeprefix(prefix).strip() 153 if reference: 154 references.add(reference) 155 if not references: 156 raise DmlRepoError("docker load did not report any image references") 157 image_ids = set() 158 for reference in sorted(references): 159 inspection = _inspect_image(reference) 160 if inspection is None: 161 raise DmlRepoError(f"docker load reported an unavailable image: {reference}") 162 image_ids.add(inspection["Id"]) 163 if len(image_ids) != 1: 164 raise DmlRepoError("docker load references must resolve to exactly one image ID") 165 return image_ids.pop() 166 167 def start( 168 self, 169 *, 170 cache_key: str, 171 execution_id: str, 172 runnable: dict[str, Any], 173 remote: dict[str, str], 174 scratch_uri: str, 175 ) -> dict[str, Any]: 176 sub = runnable.get("sub") 177 if sub is None: 178 raise DmlRepoError("docker executor requires sub runnable") 179 input_uri = _scratch_uri(scratch_uri, "input.json") 180 output_uri = _scratch_uri(scratch_uri, "output.json") 181 182 workdir = Path(tempfile.mkdtemp(prefix=f"dml-docker-{execution_id}-")) 183 payload = json.dumps( 184 { 185 "operation": "invoke", 186 "runnable": sub, 187 "cache_key": cache_key, 188 "execution_id": execution_id, 189 "remote": remote, 190 "scratch_uri": scratch_uri, 191 "adapter_state": None, 192 } 193 ) 194 try: 195 _write_scratch_json(input_uri, payload, raw=True) 196 staged_image = self._stage_image(runnable, workdir, remote) 197 with _image_lock(): 198 image_ref = self._prepare_image(staged_image) 199 container_id = self._run_docker( 200 "run", 201 "-d", 202 *cast(list[str], runnable.get("kwargs", {}).get("flags", [])), 203 "--label", 204 _CONTAINER_LABEL, 205 "-e", 206 f"DML_REMOTE_ROOT={remote['root']}", 207 image_ref, 208 sub["adapter"], 209 "--poll", 210 "-i", 211 input_uri, 212 "-o", 213 output_uri, 214 ) 215 finally: 216 shutil.rmtree(workdir, ignore_errors=True) 217 218 return { 219 "status": "retry", 220 "error": None, 221 "state": { 222 "owner": _current_owner(), 223 "container_id": container_id, 224 "image": image_ref, 225 }, 226 } 227 228 def poll( 229 self, 230 cache_key: str, 231 execution_id: str, 232 runnable: dict[str, Any], 233 state: dict[str, Any], 234 remote: dict[str, str], 235 scratch_uri: str, 236 ) -> dict[str, Any]: 237 del cache_key, execution_id, runnable, remote 238 if state["owner"] != _current_owner(): 239 return {"status": "retry", "error": None, "state": state} 240 container_id = state.get("container_id") 241 242 if not isinstance(container_id, str) or not container_id: 243 return { 244 "status": "failure", 245 "error": "docker poll: missing container_id in job state", 246 "state": state, 247 } 248 249 docker_bin = shutil.which("docker") 250 if docker_bin is None: 251 return { 252 "status": "failure", 253 "error": "docker poll: docker executable not found", 254 "state": state, 255 } 256 257 proc = subprocess.run( 258 [docker_bin, "inspect", "--format", "{{.State.Status}}", container_id], 259 check=False, 260 capture_output=True, 261 text=True, 262 ) 263 if proc.returncode != 0: 264 # An inspect error cannot establish that the container has exited. 265 return {"status": "retry", "error": None, "state": state} 266 else: 267 container_status = proc.stdout.strip() 268 269 if container_status in ("created", "running", "paused", "restarting"): 270 return {"status": "retry", "error": None, "state": state} 271 272 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 273 if raw is not None: 274 try: 275 result = validate_adapter_response(json.loads(raw)) 276 result.pop("adapter_state", None) 277 return {**result, "state": state} 278 except Exception as e: 279 raise DmlRepoError(f"docker poll: invalid nested adapter output: {e}") from e 280 281 return { 282 "status": "failure", 283 "error": f"docker container {container_id} exited without output", 284 "state": state, 285 } 286 287 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 288 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 289 if isinstance(state, dict) and state["owner"] != _current_owner(): 290 return {"status": "retry", "error": None, "state": state} 291 state = state if isinstance(state, dict) else {} 292 container_id = state.get("container_id") 293 docker_bin = shutil.which("docker") 294 if not isinstance(container_id, str) or not container_id or docker_bin is None: 295 return {"status": "success", "error": None, "state": state} 296 try: 297 status = self._run_docker( 298 "inspect", "--format", "{{.State.Status}}", container_id, docker_bin=docker_bin 299 ) 300 if status in {"created", "running", "paused", "restarting"}: 301 return {"status": "retry", "error": None, "state": state} 302 self._run_docker("update", "--restart=no", container_id, docker_bin=docker_bin) 303 except DmlRepoError as exc: 304 if "No such container:" not in str(exc) and "No such object:" not in str(exc): 305 raise 306 return {"status": "success", "error": None, "state": state} 307 308 def cancel( 309 self, 310 cache_key: str, 311 execution_id: str, 312 runnable: dict[str, Any], 313 state: dict[str, Any], 314 remote: dict[str, str], 315 scratch_uri: str, 316 cancel_requested_by: str | None, 317 argv_ref: str | None = None, 318 ) -> dict[str, Any]: 319 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 320 if isinstance(state, dict) and state["owner"] != _current_owner(): 321 return {"status": "retry", "error": None, "state": state} 322 state = state if isinstance(state, dict) else {} 323 docker_bin = shutil.which("docker") 324 container_id = state.get("container_id") 325 if docker_bin is None: 326 if isinstance(container_id, str) and container_id: 327 return {"status": "failure", "error": "docker executable not found in PATH", "state": state} 328 return {"status": "cancelled", "error": None, "state": state} 329 if isinstance(container_id, str) and container_id: 330 try: 331 self._run_docker("update", "--restart=no", container_id, docker_bin=docker_bin) 332 self._run_docker("stop", container_id, docker_bin=docker_bin) 333 except DmlRepoError as exc: 334 if "No such container:" not in str(exc): 335 raise 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}
88 @classmethod 89 def resolve_runnable(cls, uri, kwargs, sub): 90 if sub is None: 91 raise DmlRepoError("docker executor requires sub runnable") 92 image = kwargs.get("image") 93 if image is None: 94 raise DmlRepoError("docker executor requires image") 95 unknown = sorted(set(kwargs.keys()) - {"image", "flags"}) 96 if unknown: 97 raise DmlRepoError(f"Unknown docker executor kwargs: {', '.join(unknown)}") 98 return Runnable( 99 target=Uri("docker"), 100 kwargs={"image": image, "flags": kwargs.get("flags", [])}, 101 sub=sub, 102 adapter="dml-local-adapter", 103 )
167 def start( 168 self, 169 *, 170 cache_key: str, 171 execution_id: str, 172 runnable: dict[str, Any], 173 remote: dict[str, str], 174 scratch_uri: str, 175 ) -> dict[str, Any]: 176 sub = runnable.get("sub") 177 if sub is None: 178 raise DmlRepoError("docker executor requires sub runnable") 179 input_uri = _scratch_uri(scratch_uri, "input.json") 180 output_uri = _scratch_uri(scratch_uri, "output.json") 181 182 workdir = Path(tempfile.mkdtemp(prefix=f"dml-docker-{execution_id}-")) 183 payload = json.dumps( 184 { 185 "operation": "invoke", 186 "runnable": sub, 187 "cache_key": cache_key, 188 "execution_id": execution_id, 189 "remote": remote, 190 "scratch_uri": scratch_uri, 191 "adapter_state": None, 192 } 193 ) 194 try: 195 _write_scratch_json(input_uri, payload, raw=True) 196 staged_image = self._stage_image(runnable, workdir, remote) 197 with _image_lock(): 198 image_ref = self._prepare_image(staged_image) 199 container_id = self._run_docker( 200 "run", 201 "-d", 202 *cast(list[str], runnable.get("kwargs", {}).get("flags", [])), 203 "--label", 204 _CONTAINER_LABEL, 205 "-e", 206 f"DML_REMOTE_ROOT={remote['root']}", 207 image_ref, 208 sub["adapter"], 209 "--poll", 210 "-i", 211 input_uri, 212 "-o", 213 output_uri, 214 ) 215 finally: 216 shutil.rmtree(workdir, ignore_errors=True) 217 218 return { 219 "status": "retry", 220 "error": None, 221 "state": { 222 "owner": _current_owner(), 223 "container_id": container_id, 224 "image": image_ref, 225 }, 226 }
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.
228 def poll( 229 self, 230 cache_key: str, 231 execution_id: str, 232 runnable: dict[str, Any], 233 state: dict[str, Any], 234 remote: dict[str, str], 235 scratch_uri: str, 236 ) -> dict[str, Any]: 237 del cache_key, execution_id, runnable, remote 238 if state["owner"] != _current_owner(): 239 return {"status": "retry", "error": None, "state": state} 240 container_id = state.get("container_id") 241 242 if not isinstance(container_id, str) or not container_id: 243 return { 244 "status": "failure", 245 "error": "docker poll: missing container_id in job state", 246 "state": state, 247 } 248 249 docker_bin = shutil.which("docker") 250 if docker_bin is None: 251 return { 252 "status": "failure", 253 "error": "docker poll: docker executable not found", 254 "state": state, 255 } 256 257 proc = subprocess.run( 258 [docker_bin, "inspect", "--format", "{{.State.Status}}", container_id], 259 check=False, 260 capture_output=True, 261 text=True, 262 ) 263 if proc.returncode != 0: 264 # An inspect error cannot establish that the container has exited. 265 return {"status": "retry", "error": None, "state": state} 266 else: 267 container_status = proc.stdout.strip() 268 269 if container_status in ("created", "running", "paused", "restarting"): 270 return {"status": "retry", "error": None, "state": state} 271 272 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 273 if raw is not None: 274 try: 275 result = validate_adapter_response(json.loads(raw)) 276 result.pop("adapter_state", None) 277 return {**result, "state": state} 278 except Exception as e: 279 raise DmlRepoError(f"docker poll: invalid nested adapter output: {e}") from e 280 281 return { 282 "status": "failure", 283 "error": f"docker container {container_id} exited without output", 284 "state": state, 285 }
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.
287 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 288 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 289 if isinstance(state, dict) and state["owner"] != _current_owner(): 290 return {"status": "retry", "error": None, "state": state} 291 state = state if isinstance(state, dict) else {} 292 container_id = state.get("container_id") 293 docker_bin = shutil.which("docker") 294 if not isinstance(container_id, str) or not container_id or docker_bin is None: 295 return {"status": "success", "error": None, "state": state} 296 try: 297 status = self._run_docker( 298 "inspect", "--format", "{{.State.Status}}", container_id, docker_bin=docker_bin 299 ) 300 if status in {"created", "running", "paused", "restarting"}: 301 return {"status": "retry", "error": None, "state": state} 302 self._run_docker("update", "--restart=no", container_id, docker_bin=docker_bin) 303 except DmlRepoError as exc: 304 if "No such container:" not in str(exc) and "No such object:" not in str(exc): 305 raise 306 return {"status": "success", "error": None, "state": state}
Idempotently prune resources after a result was published.
308 def cancel( 309 self, 310 cache_key: str, 311 execution_id: str, 312 runnable: dict[str, Any], 313 state: dict[str, Any], 314 remote: dict[str, str], 315 scratch_uri: str, 316 cancel_requested_by: str | None, 317 argv_ref: str | None = None, 318 ) -> dict[str, Any]: 319 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 320 if isinstance(state, dict) and state["owner"] != _current_owner(): 321 return {"status": "retry", "error": None, "state": state} 322 state = state if isinstance(state, dict) else {} 323 docker_bin = shutil.which("docker") 324 container_id = state.get("container_id") 325 if docker_bin is None: 326 if isinstance(container_id, str) and container_id: 327 return {"status": "failure", "error": "docker executable not found in PATH", "state": state} 328 return {"status": "cancelled", "error": None, "state": state} 329 if isinstance(container_id, str) and container_id: 330 try: 331 self._run_docker("update", "--restart=no", container_id, docker_bin=docker_bin) 332 self._run_docker("stop", container_id, docker_bin=docker_bin) 333 except DmlRepoError as exc: 334 if "No such container:" not in str(exc): 335 raise 336 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": {"DML_REMOTE_ROOT": remote["root"]}, 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": {"DML_REMOTE_ROOT": remote["root"]}, 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 )