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}
55class DockerExecutor(ExecutorBase): 56 name = "docker" 57 adapter = "local" 58 59 @classmethod 60 def resolve_runnable(cls, uri, kwargs, sub): 61 if sub is None: 62 raise DmlRepoError("docker executor requires sub runnable") 63 image = kwargs.get("image") 64 if image is None: 65 raise DmlRepoError("docker executor requires image") 66 unknown = sorted(set(kwargs.keys()) - {"image", "flags"}) 67 if unknown: 68 raise DmlRepoError(f"Unknown docker executor kwargs: {', '.join(unknown)}") 69 return Runnable( 70 target=Uri("docker"), 71 kwargs={"image": image, "flags": kwargs.get("flags", [])}, 72 sub=sub, 73 adapter="dml-local-adapter", 74 ) 75 76 @staticmethod 77 def _run_docker(*args: str, check: bool = True, docker_bin: str | None = None) -> str: 78 docker_bin = docker_bin or shutil.which("docker") 79 if docker_bin is None: 80 raise DmlRepoError("docker executable not found in PATH") 81 proc = subprocess.run([docker_bin, *args], check=False, capture_output=True, text=True) 82 if proc.returncode == 0: 83 return proc.stdout.strip() or proc.stderr.strip() 84 if check: 85 command = f"{docker_bin} {' '.join(args)}" 86 raise DmlRepoError( 87 f"docker command failed ({proc.returncode}): {command}\nSTDOUT:\n{proc.stdout}\nSTDERR:\n{proc.stderr}" 88 ) 89 return proc.stdout.strip() or proc.stderr.strip() 90 91 @staticmethod 92 def _image_input(runnable: dict[str, Any]) -> str: 93 image = runnable.get("kwargs", {}).get("image") 94 if isinstance(image, dict): 95 image = image.get("uri") 96 if isinstance(image, Uri): 97 return image.uri 98 if isinstance(image, str) and image: 99 return image 100 raise DmlRepoError("docker executor image must resolve to a non-empty Uri or string") 101 102 @staticmethod 103 def _image_tag_from_tar(tar_path: Path) -> str: 104 with tarfile.open(tar_path, mode="r") as tf: 105 member = tf.extractfile("manifest.json") 106 if member is None: 107 raise DmlRepoError("docker image tar missing manifest.json") 108 manifest = json.loads(member.read()) 109 repo_tags = manifest[0].get("RepoTags") if manifest else None 110 if not isinstance(repo_tags, list) or not repo_tags or not isinstance(repo_tags[0], str) or not repo_tags[0]: 111 raise DmlRepoError("docker image tar missing RepoTags") 112 return cast(str, repo_tags[0]) 113 114 @staticmethod 115 def _isolate_image_tar(tar_path: Path, isolated_path: Path) -> str: 116 DockerExecutor._image_tag_from_tar(tar_path) 117 isolated_tag = f"daggerml-execution:{uuid4().hex}" 118 with tarfile.open(tar_path, mode="r:*") as source, tarfile.open(isolated_path, mode="w") as target: 119 for member in source: 120 if member.name == "repositories": 121 continue # Legacy Docker metadata would also restore the original tag. 122 content = source.extractfile(member) if member.isfile() else None 123 if member.name == "manifest.json": 124 assert content is not None 125 manifest = json.load(content) 126 if not isinstance(manifest, list) or len(manifest) != 1: 127 raise DmlRepoError("docker image tar must contain exactly one image") 128 manifest[0]["RepoTags"] = [isolated_tag] 129 data = json.dumps(manifest).encode("utf-8") 130 member = copy.copy(member) 131 member.size = len(data) 132 content = io.BytesIO(data) 133 target.addfile(member, content) 134 return isolated_tag 135 136 @staticmethod 137 def _prepare_image(runnable: dict[str, Any], workdir: Path, remote: dict[str, Any]) -> tuple[str, str | None]: 138 image = DockerExecutor._image_input(runnable) 139 if not is_s3_uri(image): 140 return image, None 141 tar_path = workdir / "image.tar" 142 store = S3Store.from_remote_root(cast(str, remote["root"])) 143 tar_path.write_bytes(store.get(image)) 144 isolated_path = workdir / "isolated.tar" 145 image_ref = DockerExecutor._isolate_image_tar(tar_path, isolated_path) 146 DockerExecutor._run_docker("load", "-i", str(isolated_path)) 147 return image_ref, image_ref 148 149 def start( 150 self, 151 *, 152 cache_key: str, 153 execution_id: str, 154 runnable: dict[str, Any], 155 remote: dict[str, str], 156 scratch_uri: str, 157 ) -> dict[str, Any]: 158 sub = runnable.get("sub") 159 if sub is None: 160 raise DmlRepoError("docker executor requires sub runnable") 161 input_uri = _scratch_uri(scratch_uri, "input.json") 162 output_uri = _scratch_uri(scratch_uri, "output.json") 163 164 workdir = Path(tempfile.mkdtemp(prefix=f"dml-docker-{execution_id}-")) 165 try: 166 image_ref, cleanup_image = self._prepare_image(runnable, workdir, remote) 167 finally: 168 shutil.rmtree(workdir, ignore_errors=True) 169 170 payload = json.dumps( 171 { 172 "operation": "invoke", 173 "runnable": sub, 174 "cache_key": cache_key, 175 "execution_id": execution_id, 176 "remote": remote, 177 "scratch_uri": scratch_uri, 178 "adapter_state": None, 179 } 180 ) 181 _write_scratch_json(input_uri, payload, raw=True) 182 183 container_id = self._run_docker( 184 "run", 185 "-d", 186 *cast(list[str], runnable.get("kwargs", {}).get("flags", [])), 187 "-e", 188 f"DML_REMOTE_ROOT={remote['root']}", 189 image_ref, 190 sub["adapter"], 191 "--poll", 192 "-i", 193 input_uri, 194 "-o", 195 output_uri, 196 ) 197 198 return { 199 "status": "retry", 200 "error": None, 201 "state": { 202 "container_id": container_id, 203 "cleanup_image": cleanup_image, 204 }, 205 } 206 207 def poll( 208 self, 209 cache_key: str, 210 execution_id: str, 211 runnable: dict[str, Any], 212 state: dict[str, Any], 213 remote: dict[str, str], 214 scratch_uri: str, 215 ) -> dict[str, Any]: 216 del cache_key, execution_id, runnable, remote 217 container_id = state.get("container_id") 218 219 if not isinstance(container_id, str) or not container_id: 220 return { 221 "status": "failure", 222 "error": "docker poll: missing container_id in job state", 223 "state": None, 224 } 225 226 docker_bin = shutil.which("docker") 227 if docker_bin is None: 228 return { 229 "status": "failure", 230 "error": "docker poll: docker executable not found", 231 "state": None, 232 } 233 234 proc = subprocess.run( 235 [docker_bin, "inspect", "--format", "{{.State.Status}}", container_id], 236 check=False, 237 capture_output=True, 238 text=True, 239 ) 240 if proc.returncode != 0: 241 # An inspect error cannot establish that the container has exited. 242 return {"status": "retry", "error": None, "state": state} 243 else: 244 container_status = proc.stdout.strip() 245 246 if container_status in ("created", "running", "paused", "restarting"): 247 return {"status": "retry", "error": None, "state": state} 248 249 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 250 if raw is not None: 251 try: 252 result = validate_adapter_response(json.loads(raw)) 253 return result 254 except Exception as e: 255 raise DmlRepoError(f"docker poll: invalid nested adapter output: {e}") from e 256 257 return { 258 "status": "failure", 259 "error": f"docker container {container_id} exited without output", 260 "state": None, 261 } 262 263 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 264 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 265 state = state if isinstance(state, dict) else {} 266 container_id = state.get("container_id") 267 docker_bin = shutil.which("docker") 268 if not isinstance(container_id, str) or not container_id or docker_bin is None: 269 return {"status": "success", "error": None, "state": state} 270 status = self._run_docker( 271 "inspect", "--format", "{{.State.Status}}", container_id, check=False, docker_bin=docker_bin 272 ) 273 if status in {"created", "running", "paused", "restarting"}: 274 return {"status": "retry", "error": None, "state": state} 275 _cleanup_docker(container_id, state.get("cleanup_image"), docker_bin) 276 return {"status": "success", "error": None, "state": state} 277 278 def cancel( 279 self, 280 cache_key: str, 281 execution_id: str, 282 runnable: dict[str, Any], 283 state: dict[str, Any], 284 remote: dict[str, str], 285 scratch_uri: str, 286 cancel_requested_by: str | None, 287 argv_ref: str | None = None, 288 ) -> dict[str, Any]: 289 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 290 state = state if isinstance(state, dict) else {} 291 docker_bin = shutil.which("docker") 292 container_id = state.get("container_id") 293 if docker_bin is None: 294 if isinstance(container_id, str) and container_id: 295 return {"status": "failure", "error": "docker executable not found in PATH", "state": state} 296 return {"status": "cancelled", "error": None, "state": state} 297 if isinstance(container_id, str) and container_id: 298 _cleanup_docker(container_id, state.get("cleanup_image"), docker_bin) 299 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}
59 @classmethod 60 def resolve_runnable(cls, uri, kwargs, sub): 61 if sub is None: 62 raise DmlRepoError("docker executor requires sub runnable") 63 image = kwargs.get("image") 64 if image is None: 65 raise DmlRepoError("docker executor requires image") 66 unknown = sorted(set(kwargs.keys()) - {"image", "flags"}) 67 if unknown: 68 raise DmlRepoError(f"Unknown docker executor kwargs: {', '.join(unknown)}") 69 return Runnable( 70 target=Uri("docker"), 71 kwargs={"image": image, "flags": kwargs.get("flags", [])}, 72 sub=sub, 73 adapter="dml-local-adapter", 74 )
149 def start( 150 self, 151 *, 152 cache_key: str, 153 execution_id: str, 154 runnable: dict[str, Any], 155 remote: dict[str, str], 156 scratch_uri: str, 157 ) -> dict[str, Any]: 158 sub = runnable.get("sub") 159 if sub is None: 160 raise DmlRepoError("docker executor requires sub runnable") 161 input_uri = _scratch_uri(scratch_uri, "input.json") 162 output_uri = _scratch_uri(scratch_uri, "output.json") 163 164 workdir = Path(tempfile.mkdtemp(prefix=f"dml-docker-{execution_id}-")) 165 try: 166 image_ref, cleanup_image = self._prepare_image(runnable, workdir, remote) 167 finally: 168 shutil.rmtree(workdir, ignore_errors=True) 169 170 payload = json.dumps( 171 { 172 "operation": "invoke", 173 "runnable": sub, 174 "cache_key": cache_key, 175 "execution_id": execution_id, 176 "remote": remote, 177 "scratch_uri": scratch_uri, 178 "adapter_state": None, 179 } 180 ) 181 _write_scratch_json(input_uri, payload, raw=True) 182 183 container_id = self._run_docker( 184 "run", 185 "-d", 186 *cast(list[str], runnable.get("kwargs", {}).get("flags", [])), 187 "-e", 188 f"DML_REMOTE_ROOT={remote['root']}", 189 image_ref, 190 sub["adapter"], 191 "--poll", 192 "-i", 193 input_uri, 194 "-o", 195 output_uri, 196 ) 197 198 return { 199 "status": "retry", 200 "error": None, 201 "state": { 202 "container_id": container_id, 203 "cleanup_image": cleanup_image, 204 }, 205 }
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.
207 def poll( 208 self, 209 cache_key: str, 210 execution_id: str, 211 runnable: dict[str, Any], 212 state: dict[str, Any], 213 remote: dict[str, str], 214 scratch_uri: str, 215 ) -> dict[str, Any]: 216 del cache_key, execution_id, runnable, remote 217 container_id = state.get("container_id") 218 219 if not isinstance(container_id, str) or not container_id: 220 return { 221 "status": "failure", 222 "error": "docker poll: missing container_id in job state", 223 "state": None, 224 } 225 226 docker_bin = shutil.which("docker") 227 if docker_bin is None: 228 return { 229 "status": "failure", 230 "error": "docker poll: docker executable not found", 231 "state": None, 232 } 233 234 proc = subprocess.run( 235 [docker_bin, "inspect", "--format", "{{.State.Status}}", container_id], 236 check=False, 237 capture_output=True, 238 text=True, 239 ) 240 if proc.returncode != 0: 241 # An inspect error cannot establish that the container has exited. 242 return {"status": "retry", "error": None, "state": state} 243 else: 244 container_status = proc.stdout.strip() 245 246 if container_status in ("created", "running", "paused", "restarting"): 247 return {"status": "retry", "error": None, "state": state} 248 249 raw = _read_scratch_output(_scratch_uri(scratch_uri, "output.json")) 250 if raw is not None: 251 try: 252 result = validate_adapter_response(json.loads(raw)) 253 return result 254 except Exception as e: 255 raise DmlRepoError(f"docker poll: invalid nested adapter output: {e}") from e 256 257 return { 258 "status": "failure", 259 "error": f"docker container {container_id} exited without output", 260 "state": None, 261 }
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.
263 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 264 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 265 state = state if isinstance(state, dict) else {} 266 container_id = state.get("container_id") 267 docker_bin = shutil.which("docker") 268 if not isinstance(container_id, str) or not container_id or docker_bin is None: 269 return {"status": "success", "error": None, "state": state} 270 status = self._run_docker( 271 "inspect", "--format", "{{.State.Status}}", container_id, check=False, docker_bin=docker_bin 272 ) 273 if status in {"created", "running", "paused", "restarting"}: 274 return {"status": "retry", "error": None, "state": state} 275 _cleanup_docker(container_id, state.get("cleanup_image"), docker_bin) 276 return {"status": "success", "error": None, "state": state}
Idempotently prune resources after a result was published.
278 def cancel( 279 self, 280 cache_key: str, 281 execution_id: str, 282 runnable: dict[str, Any], 283 state: dict[str, Any], 284 remote: dict[str, str], 285 scratch_uri: str, 286 cancel_requested_by: str | None, 287 argv_ref: str | None = None, 288 ) -> dict[str, Any]: 289 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 290 state = state if isinstance(state, dict) else {} 291 docker_bin = shutil.which("docker") 292 container_id = state.get("container_id") 293 if docker_bin is None: 294 if isinstance(container_id, str) and container_id: 295 return {"status": "failure", "error": "docker executable not found in PATH", "state": state} 296 return {"status": "cancelled", "error": None, "state": state} 297 if isinstance(container_id, str) and container_id: 298 _cleanup_docker(container_id, state.get("cleanup_image"), docker_bin) 299 return {"status": "cancelled", "error": None, "state": state}
32class ScriptExecutor(ExecutorBase): 33 name = "script" 34 adapter = "local" 35 36 ############################# resolve runnable ############################# 37 ############################################################################ 38 39 @classmethod 40 def resolve_runnable(cls, uri, kwargs, sub): 41 if sub is not None: 42 raise DmlRepoError("script executor does not accept sub runnable") 43 resolved_kw, script = cls._script_kwargs(dict(kwargs)) 44 resolved_kw["script_uri"] = S3Store().put(data=script.encode("utf-8"), suffix=".py") 45 return Runnable(target=Uri("script"), kwargs=resolved_kw, sub=sub, adapter="dml-local-adapter") 46 47 @classmethod 48 def _script_kwargs(cls, kwargs: dict) -> tuple[dict, str]: 49 allowed = {"fn", "fn_name", "script", "prepop", "extra_objs", "post_lines", "tags"} 50 unknown = sorted(set(kwargs.keys()) - allowed) 51 if unknown: 52 bad = ", ".join(unknown) 53 raise DmlRepoError(f"Unknown script executor kwargs: {bad}") 54 fn = kwargs.get("fn") 55 prepop = kwargs.get("prepop", {}) 56 extra_objs = list(kwargs.get("extra_objs", [])) 57 post_lines = list(kwargs.get("post_lines", [])) 58 tags = kwargs.get("tags", []) 59 if not isinstance(tags, list) or not all(isinstance(tag, str) for tag in tags): 60 raise DmlRepoError("script tags must be a list of strings") 61 script = kwargs.get("script") 62 fn_name = kwargs.get("fn_name") 63 if (script is None) != (fn_name is None): 64 raise DmlRepoError("script and fn_name must be captured together") 65 if script is None: 66 if not callable(fn): 67 raise DmlRepoError("script fn must be callable") 68 params = list(inspect.signature(fn).parameters.values()) 69 if not params: 70 raise DmlRepoError("script fn must include at least one parameter") 71 script = cls._render_script(fn, extra_objs=extra_objs, post_lines=post_lines) 72 fn_name = fn.__name__ 73 elif not isinstance(script, str) or not isinstance(fn_name, str): 74 raise DmlRepoError("captured script and fn_name must be strings") 75 resolved = {"prepop": prepop, "fn_name": fn_name} 76 if tags: 77 resolved["tags"] = sorted(set(tags)) 78 return resolved, script 79 80 @classmethod 81 def _render_script(cls, fn, extra_objs: list, post_lines: list[str]) -> str: 82 chunks: list[str] = [] 83 for obj in [*extra_objs, fn]: 84 try: 85 raw = dedent(inspect.getsource(inspect.unwrap(obj))).strip() 86 chunks.append(cls._strip_funkify_decorators(raw)) 87 except (OSError, TypeError) as e: 88 raise DmlRepoError(f"Failed to serialize object source: {e}") from e 89 if post_lines: 90 chunks.extend(post_lines) 91 script = "\n\n".join(chunks) 92 try: 93 mod = ast.parse(script) 94 except SyntaxError as e: 95 raise DmlRepoError(f"Generated script is not valid Python: {e}") from e 96 if not any(isinstance(n, ast.FunctionDef) and n.name == fn.__name__ for n in mod.body): 97 raise DmlRepoError(f"Function '{fn.__name__}' is not globally defined in generated script") 98 return script 99 100 @staticmethod 101 def _strip_funkify_decorators(source: str) -> str: 102 module = ast.parse(source) 103 for node in module.body: 104 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): 105 node.decorator_list = [] 106 return ast.unparse(module).strip() 107 108 ############################# start/poll/cancel ############################ 109 110 def start( 111 self, 112 cache_key: str, 113 execution_id: str, 114 runnable: dict[str, Any], 115 remote: dict[str, str], 116 scratch_uri: str, 117 ) -> AdapterInvokeResponse: 118 del scratch_uri 119 tags = runnable.get("kwargs", {}).get("tags") 120 workdir = Path(tempfile.mkdtemp(prefix=f"dml-script-{execution_id}-")) 121 payload_path = workdir / "supervisor-input.json" 122 result_path = workdir / "result.json" 123 stdout_path = workdir / "stdout.log" 124 stderr_path = workdir / "stderr.log" 125 cmd = [ 126 sys.executable, 127 "-m", 128 "daggerml.contrib.executors.script", 129 "--execution-id", 130 execution_id, 131 "--cache-key", 132 cache_key, 133 "--remote-root", 134 remote["root"], 135 ] 136 if tags: 137 cmd.extend(["--tags", json.dumps(tags, separators=(",", ":"))]) 138 payload = { 139 "version": 0, 140 "cache_key": cache_key, 141 "execution_id": execution_id, 142 "cmd": cmd, 143 "remote": remote, 144 "env": {}, 145 } 146 payload_path.write_text(json.dumps(payload, separators=(",", ":"), sort_keys=True)) 147 with stdout_path.open("w") as stdout_f, stderr_path.open("w") as stderr_f: 148 proc = subprocess.Popen( 149 [ 150 sys.executable, 151 "-m", 152 "daggerml.contrib.supervisor", 153 "-i", 154 str(payload_path), 155 "-o", 156 str(result_path), 157 ], 158 stdout=stdout_f, 159 stderr=stderr_f, 160 start_new_session=True, 161 close_fds=True, 162 env={**os.environ, "PYTHONUNBUFFERED": "1", **payload["env"]}, 163 ) 164 launch_state = { 165 "pid": proc.pid, 166 "workdir": str(workdir), 167 "result_path": str(result_path), 168 "stdout_path": str(stdout_path), 169 "stderr_path": str(stderr_path), 170 } 171 return {"status": "retry", "error": None, "state": launch_state} 172 173 def poll( 174 self, 175 cache_key: str, 176 execution_id: str, 177 runnable: dict[str, Any], 178 state: dict[str, Any], 179 remote: dict[str, str], 180 scratch_uri: str, 181 ) -> AdapterInvokeResponse: 182 del cache_key, execution_id, runnable, remote, scratch_uri 183 result_path = Path(state.get("result_path", "")) 184 pid = state.get("pid") 185 # Polls may run either in the launching adapter process or in a later 186 # process. Reap children when we can; otherwise fall back to a direct 187 # PID probe for cross-process polling. 188 if isinstance(pid, int): 189 try: 190 done_pid, _ = os.waitpid(pid, os.WNOHANG) 191 if done_pid == 0: 192 return {"status": "retry", "error": None, "state": state} 193 except ChildProcessError: 194 try: 195 os.kill(pid, 0) 196 return {"status": "retry", "error": None, "state": state} 197 except ProcessLookupError: 198 pass 199 except PermissionError: 200 return {"status": "retry", "error": None, "state": state} 201 # Process exited — read result 202 if result_path.exists(): 203 try: 204 parsed = json.loads(result_path.read_text()) 205 if parsed.get("status") in {"succeeded", "failed"}: 206 return { 207 "status": "success" if parsed["status"] == "succeeded" else "failure", 208 "error": parsed.get("error"), 209 "state": state, 210 } 211 except Exception as e: 212 return { 213 "status": "failure", 214 "error": f"Could not read supervisor result: {e}", 215 "state": state, 216 } 217 return { 218 "status": "failure", 219 "error": "Script supervisor exited without result", 220 "state": state, 221 } 222 223 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 224 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 225 state = state if isinstance(state, dict) else {} 226 pid = state.get("pid") 227 if isinstance(pid, int): 228 try: 229 done_pid, _ = os.waitpid(pid, os.WNOHANG) 230 active = done_pid == 0 231 except ChildProcessError: 232 try: 233 os.kill(pid, 0) 234 active = True 235 except ProcessLookupError: 236 active = False 237 except PermissionError as exc: 238 return {"status": "failure", "error": f"script cleanup failed: {exc}", "state": state} 239 if active: 240 # The DML result is already published, so the supervisor is no longer needed. 241 # Stop its process group before removing scratch so it cannot keep writing there. 242 try: 243 os.killpg(pid, signal.SIGTERM) 244 except ProcessLookupError: 245 pass 246 except PermissionError as exc: 247 return {"status": "failure", "error": f"script cleanup failed: {exc}", "state": state} 248 _cleanup_workdir(state) 249 workdir = state.get("workdir") 250 if isinstance(workdir, str) and os.path.exists(workdir): 251 return {"status": "failure", "error": f"script cleanup failed to remove {workdir}", "state": state} 252 return {"status": "success", "error": None, "state": state} 253 254 def cancel( 255 self, 256 cache_key: str, 257 execution_id: str, 258 runnable: dict[str, Any], 259 state: dict[str, Any], 260 remote: dict[str, str], 261 scratch_uri: str, 262 cancel_requested_by: str | None, 263 argv_ref: str | None = None, 264 ) -> AdapterCancelResponse: 265 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 266 if not isinstance(state, dict): 267 return {"status": "cancelled", "error": None, "state": {}} 268 pid = state.get("pid") 269 if isinstance(pid, int): 270 try: 271 os.killpg(pid, signal.SIGTERM) 272 except ProcessLookupError: 273 pass 274 except PermissionError as exc: 275 return {"status": "failure", "error": f"script cancellation failed: {exc}", "state": state} 276 _cleanup_workdir(state) 277 workdir = state.get("workdir") 278 if isinstance(workdir, str) and os.path.exists(workdir): 279 return {"status": "failure", "error": f"script cancellation failed to remove {workdir}", "state": state} 280 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}
39 @classmethod 40 def resolve_runnable(cls, uri, kwargs, sub): 41 if sub is not None: 42 raise DmlRepoError("script executor does not accept sub runnable") 43 resolved_kw, script = cls._script_kwargs(dict(kwargs)) 44 resolved_kw["script_uri"] = S3Store().put(data=script.encode("utf-8"), suffix=".py") 45 return Runnable(target=Uri("script"), kwargs=resolved_kw, sub=sub, adapter="dml-local-adapter")
110 def start( 111 self, 112 cache_key: str, 113 execution_id: str, 114 runnable: dict[str, Any], 115 remote: dict[str, str], 116 scratch_uri: str, 117 ) -> AdapterInvokeResponse: 118 del scratch_uri 119 tags = runnable.get("kwargs", {}).get("tags") 120 workdir = Path(tempfile.mkdtemp(prefix=f"dml-script-{execution_id}-")) 121 payload_path = workdir / "supervisor-input.json" 122 result_path = workdir / "result.json" 123 stdout_path = workdir / "stdout.log" 124 stderr_path = workdir / "stderr.log" 125 cmd = [ 126 sys.executable, 127 "-m", 128 "daggerml.contrib.executors.script", 129 "--execution-id", 130 execution_id, 131 "--cache-key", 132 cache_key, 133 "--remote-root", 134 remote["root"], 135 ] 136 if tags: 137 cmd.extend(["--tags", json.dumps(tags, separators=(",", ":"))]) 138 payload = { 139 "version": 0, 140 "cache_key": cache_key, 141 "execution_id": execution_id, 142 "cmd": cmd, 143 "remote": remote, 144 "env": {}, 145 } 146 payload_path.write_text(json.dumps(payload, separators=(",", ":"), sort_keys=True)) 147 with stdout_path.open("w") as stdout_f, stderr_path.open("w") as stderr_f: 148 proc = subprocess.Popen( 149 [ 150 sys.executable, 151 "-m", 152 "daggerml.contrib.supervisor", 153 "-i", 154 str(payload_path), 155 "-o", 156 str(result_path), 157 ], 158 stdout=stdout_f, 159 stderr=stderr_f, 160 start_new_session=True, 161 close_fds=True, 162 env={**os.environ, "PYTHONUNBUFFERED": "1", **payload["env"]}, 163 ) 164 launch_state = { 165 "pid": proc.pid, 166 "workdir": str(workdir), 167 "result_path": str(result_path), 168 "stdout_path": str(stdout_path), 169 "stderr_path": str(stderr_path), 170 } 171 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.
173 def poll( 174 self, 175 cache_key: str, 176 execution_id: str, 177 runnable: dict[str, Any], 178 state: dict[str, Any], 179 remote: dict[str, str], 180 scratch_uri: str, 181 ) -> AdapterInvokeResponse: 182 del cache_key, execution_id, runnable, remote, scratch_uri 183 result_path = Path(state.get("result_path", "")) 184 pid = state.get("pid") 185 # Polls may run either in the launching adapter process or in a later 186 # process. Reap children when we can; otherwise fall back to a direct 187 # PID probe for cross-process polling. 188 if isinstance(pid, int): 189 try: 190 done_pid, _ = os.waitpid(pid, os.WNOHANG) 191 if done_pid == 0: 192 return {"status": "retry", "error": None, "state": state} 193 except ChildProcessError: 194 try: 195 os.kill(pid, 0) 196 return {"status": "retry", "error": None, "state": state} 197 except ProcessLookupError: 198 pass 199 except PermissionError: 200 return {"status": "retry", "error": None, "state": state} 201 # Process exited — read result 202 if result_path.exists(): 203 try: 204 parsed = json.loads(result_path.read_text()) 205 if parsed.get("status") in {"succeeded", "failed"}: 206 return { 207 "status": "success" if parsed["status"] == "succeeded" else "failure", 208 "error": parsed.get("error"), 209 "state": state, 210 } 211 except Exception as e: 212 return { 213 "status": "failure", 214 "error": f"Could not read supervisor result: {e}", 215 "state": state, 216 } 217 return { 218 "status": "failure", 219 "error": "Script supervisor exited without result", 220 "state": state, 221 }
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.
223 def cleanup(self, cache_key, execution_id, runnable, state, remote, scratch_uri, result_ref): 224 del cache_key, execution_id, runnable, remote, scratch_uri, result_ref 225 state = state if isinstance(state, dict) else {} 226 pid = state.get("pid") 227 if isinstance(pid, int): 228 try: 229 done_pid, _ = os.waitpid(pid, os.WNOHANG) 230 active = done_pid == 0 231 except ChildProcessError: 232 try: 233 os.kill(pid, 0) 234 active = True 235 except ProcessLookupError: 236 active = False 237 except PermissionError as exc: 238 return {"status": "failure", "error": f"script cleanup failed: {exc}", "state": state} 239 if active: 240 # The DML result is already published, so the supervisor is no longer needed. 241 # Stop its process group before removing scratch so it cannot keep writing there. 242 try: 243 os.killpg(pid, signal.SIGTERM) 244 except ProcessLookupError: 245 pass 246 except PermissionError as exc: 247 return {"status": "failure", "error": f"script cleanup failed: {exc}", "state": state} 248 _cleanup_workdir(state) 249 workdir = state.get("workdir") 250 if isinstance(workdir, str) and os.path.exists(workdir): 251 return {"status": "failure", "error": f"script cleanup failed to remove {workdir}", "state": state} 252 return {"status": "success", "error": None, "state": state}
Idempotently prune resources after a result was published.
254 def cancel( 255 self, 256 cache_key: str, 257 execution_id: str, 258 runnable: dict[str, Any], 259 state: dict[str, Any], 260 remote: dict[str, str], 261 scratch_uri: str, 262 cancel_requested_by: str | None, 263 argv_ref: str | None = None, 264 ) -> AdapterCancelResponse: 265 del cache_key, execution_id, runnable, remote, scratch_uri, cancel_requested_by, argv_ref 266 if not isinstance(state, dict): 267 return {"status": "cancelled", "error": None, "state": {}} 268 pid = state.get("pid") 269 if isinstance(pid, int): 270 try: 271 os.killpg(pid, signal.SIGTERM) 272 except ProcessLookupError: 273 pass 274 except PermissionError as exc: 275 return {"status": "failure", "error": f"script cancellation failed: {exc}", "state": state} 276 _cleanup_workdir(state) 277 workdir = state.get("workdir") 278 if isinstance(workdir, str) and os.path.exists(workdir): 279 return {"status": "failure", "error": f"script cancellation failed to remove {workdir}", "state": state} 280 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 )