hpcman.queue.acctinfo¶

class QueueAcctinfo(BaseModel):
    """QueueAcct class based on Pydantic BaseModel.

    Attributes:
        users list[str]: User names.
        partitions list[str]: Queue names.
        accounts list[str]: Account names.
        debug (bool, optional): Debug flag. Defaults to False.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)
    users: list[str] = []
    partitions: list[str] = []
    accounts: list[str] = []
    debug: bool

    priority_partitions: set[str] = set()
    partition_map: dict[str, set[str]] = dict()
    assoc: SlurmAssociations | None = None

    @field_validator("users")  # type: ignore
    @classmethod
    def check_users(cls, v: list[str]) -> list[str]:
        if not any([check_valid_slurm_association(x) for x in v]):
            raise InvalidSlurmUser(f"One or more provided users is invalid: {v}")
        else:
            return v

    @field_validator("accounts")  # type: ignore
    @classmethod
    def check_accounts(cls, v: list[str]) -> list[str]:
        if not any([check_valid_slurm_association(x, "accounts") for x in v]):
            raise InvalidSlurmAccount(f"One or more provided accounts is invalid: {v}")
        else:
            return v

    @field_validator("partitions")  # type: ignore
    @classmethod
    def check_partitions(cls, v: list[str]) -> list[str]:
        if not any([check_valid_slurm_partition(part) for part in v]):
            raise InvalidSlurmPartition(f"One or more provided partitions is invalid: {v}")
        else:
            return v

    def generate_acct_tree(self) -> Tree:
        user_tree = Tree("Slurm Queue Account Associations")
        filtered_users = _gather_acctinfo_data(self)

        for fu in filtered_users:
            user_text = Text()
            user_text.append(fu.name, style="cyan")
            user_text.append(" Default Acct: ")
            user_text.append(fu.default_account, style="green")
            user_info = Tree(user_text)
            for fa in fu.accounts:
                acct_tree = Tree(Text(fa.name))
                acct_text = Text()
                if fa.partitions:
                    for i, part in enumerate(fa.partitions):
                        if i > 0:
                            acct_text.append(",")
                        if part in self.priority_partitions:
                            acct_text.append(f"{part}*", style="bold magenta")
                        else:
                            acct_text.append(part, style="magenta")
                    acct_tree.add(Tree(acct_text))
                else:
                    acct_tree.add(Tree(Text("None", style="red")))
                user_info.add(acct_tree)
            user_tree.add(user_info)

        return user_tree

    @field_validator("users")  # type: ignore
    @classmethod
    def check_users(cls, v: list[str]) -> list[str]:
        if not any([check_valid_slurm_association(x) for x in v]):
            raise InvalidSlurmUser(f"One or more provided users is invalid: {v}")
        else:
            return v

    @field_validator("accounts")  # type: ignore
    @classmethod
    def check_accounts(cls, v: list[str]) -> list[str]:
        if not any([check_valid_slurm_association(x, "accounts") for x in v]):
            raise InvalidSlurmAccount(f"One or more provided accounts is invalid: {v}")
        else:
            return v

    @field_validator("partitions")  # type: ignore
    @classmethod
    def check_partitions(cls, v: list[str]) -> list[str]:
        if not any([check_valid_slurm_partition(part) for part in v]):
            raise InvalidSlurmPartition(f"One or more provided partitions is invalid: {v}")
        else:
            return v

    def generate_acct_tree(self) -> Tree:
        user_tree = Tree("Slurm Queue Account Associations")
        filtered_users = _gather_acctinfo_data(self)

        for fu in filtered_users:
            user_text = Text()
            user_text.append(fu.name, style="cyan")
            user_text.append(" Default Acct: ")
            user_text.append(fu.default_account, style="green")
            user_info = Tree(user_text)
            for fa in fu.accounts:
                acct_tree = Tree(Text(fa.name))
                acct_text = Text()
                if fa.partitions:
                    for i, part in enumerate(fa.partitions):
                        if i > 0:
                            acct_text.append(",")
                        if part in self.priority_partitions:
                            acct_text.append(f"{part}*", style="bold magenta")
                        else:
                            acct_text.append(part, style="magenta")
                    acct_tree.add(Tree(acct_text))
                else:
                    acct_tree.add(Tree(Text("None", style="red")))
                user_info.add(acct_tree)
            user_tree.add(user_info)

        return user_tree

@dataclass
class FilteredAccount:
    name: str
    partitions: list[str] = field(default_factory=list)
    priority_partitions: list[str] = field(default_factory=list)

@dataclass
class FilteredUser:
    name: str
    default_account: str
    accounts: list[FilteredAccount] = field(default_factory=list)

def run_queue_acctinfo(**kwargs) -> None:
    """Run the queue_acct command.

    Args:
        user (str): User name.
        queue (Optional[str], optional): Queue name. Defaults to None.
        debug (bool, optional): Debug flag. Defaults to False.

    Returns:
        None
    """
    json_output = kwargs.pop("json_output", False)
    console = Console()
    try:
        acct = QueueAcctinfo(**{k: v for k, v in kwargs.items() if v is not None})
    except ValidationError as e:
        handle_validation_errors(e, print_error=False)
    if json_output:
        _print_acctinfo_json(acct)
        return
    console.print(acct.generate_acct_tree())
    if kwargs["debug"]:
        console.print(acct)