|
21 | 21 | from datetime import datetime, timezone
|
22 | 22 | import enum
|
23 | 23 | import json
|
| 24 | +from typing import TYPE_CHECKING |
24 | 25 |
|
25 |
| -from aiida.common import AIIDA_LOGGER |
| 26 | +from aiida.common import AIIDA_LOGGER, CodeRunMode |
26 | 27 | from aiida.common.extendeddicts import AttributeDict, DefaultFieldsAttributeDict
|
27 | 28 | from aiida.common.timezone import make_aware, timezone_from_name
|
28 | 29 |
|
@@ -108,6 +109,12 @@ class NodeNumberJobResource(JobResource):
|
108 | 109 | 'num_cores_per_mpiproc',
|
109 | 110 | )
|
110 | 111 |
|
| 112 | + if TYPE_CHECKING: |
| 113 | + num_machines: int |
| 114 | + num_mpiprocs_per_machine: int |
| 115 | + num_cores_per_machine: int |
| 116 | + num_cores_per_mpiproc: int |
| 117 | + |
111 | 118 | @classmethod
|
112 | 119 | def validate_resources(cls, **kwargs):
|
113 | 120 | """Validate the resources against the job resource class of this scheduler.
|
@@ -193,6 +200,10 @@ class ParEnvJobResource(JobResource):
|
193 | 200 | 'tot_num_mpiprocs',
|
194 | 201 | )
|
195 | 202 |
|
| 203 | + if TYPE_CHECKING: |
| 204 | + parallel_env: str |
| 205 | + tot_num_mpiprocs: int |
| 206 | + |
196 | 207 | @classmethod
|
197 | 208 | def validate_resources(cls, **kwargs):
|
198 | 209 | """Validate the resources against the job resource class of this scheduler.
|
@@ -366,6 +377,34 @@ class JobTemplate(DefaultFieldsAttributeDict): # pylint: disable=too-many-insta
|
366 | 377 | 'codes_info',
|
367 | 378 | )
|
368 | 379 |
|
| 380 | + if TYPE_CHECKING: |
| 381 | + shebang: str |
| 382 | + submit_as_hold: bool |
| 383 | + rerunnable: bool |
| 384 | + job_environment: dict[str, str] | None |
| 385 | + environment_variables_double_quotes: bool | None |
| 386 | + working_directory: str |
| 387 | + email: str |
| 388 | + email_on_started: bool |
| 389 | + email_on_terminated: bool |
| 390 | + job_name: str |
| 391 | + sched_output_path: str |
| 392 | + sched_error_path: str |
| 393 | + sched_join_files: bool |
| 394 | + queue_name: str |
| 395 | + account: str |
| 396 | + qos: str |
| 397 | + job_resource: JobResource |
| 398 | + priority: str |
| 399 | + max_memory_kb: int | None |
| 400 | + max_wallclock_seconds: int |
| 401 | + custom_scheduler_commands: str |
| 402 | + prepend_text: str |
| 403 | + append_text: str |
| 404 | + import_sys_environment: bool | None |
| 405 | + codes_run_mode: CodeRunMode |
| 406 | + codes_info: list[JobTemplateCodeInfo] |
| 407 | + |
369 | 408 |
|
370 | 409 | @dataclass
|
371 | 410 | class JobTemplateCodeInfo:
|
@@ -474,6 +513,29 @@ class JobInfo(DefaultFieldsAttributeDict): # pylint: disable=too-many-instance-
|
474 | 513 | 'finish_time'
|
475 | 514 | )
|
476 | 515 |
|
| 516 | + if TYPE_CHECKING: |
| 517 | + job_id: str |
| 518 | + title: str |
| 519 | + exit_status: int |
| 520 | + terminating_signal: int |
| 521 | + annotation: str |
| 522 | + job_state: JobState |
| 523 | + job_substate: str |
| 524 | + allocated_machines: list[MachineInfo] |
| 525 | + job_owner: str |
| 526 | + num_mpiprocs: int |
| 527 | + num_cpus: int |
| 528 | + num_machines: int |
| 529 | + queue_name: str |
| 530 | + account: str |
| 531 | + qos: str |
| 532 | + wallclock_time_seconds: int |
| 533 | + requested_wallclock_time_seconds: int |
| 534 | + cpu_time: int |
| 535 | + submission_time: datetime |
| 536 | + dispatch_time: datetime |
| 537 | + finish_time: datetime |
| 538 | + |
477 | 539 | # If some fields require special serializers, specify them here.
|
478 | 540 | # You then need to define also the respective _serialize_FIELDTYPE and
|
479 | 541 | # _deserialize_FIELDTYPE methods
|
|
0 commit comments