Controlling cyclers¶

class ServerManager:
    """The ServerManager class manages the cycling servers."""

    def __init__(self) -> None:
        """Initialize the server manager object."""
        self.config = config.get_config()
        if not self.config.get("Snapshots folder path"):
            msg = "'Snapshots folder path' not found in config file. Cannot save snapshots."
            raise ValueError(msg)
        if not self.config.get("Servers"):
            logger.info("There are no servers in the configuration.")
        else:
            logger.info("Server manager initialised, consider updating database with update_db()")

    @cached_property
    def servers(self) -> dict[str, CyclerServer]:
        """Get a dictionary of Cycler Servers."""
        return _get_servers()

    def update_db(self) -> None:
        """Query all cycler servers and update the pipelines table in the database with their current status."""
        logger.info("Querying all cycler servers...")

        def fetch(label: str, server: CyclerServer) -> tuple[str, list | None]:
            t0 = time()
            try:
                pipelines = server.get_pipelines()
            except Exception:
                logger.exception("Error getting pipeline status from %s", label)
                return label, None
            logger.info("Queried server '%s' in %s s", label, time() - t0)
            return label, pipelines

        with ThreadPoolExecutor() as executor:
            futures = {executor.submit(fetch, label, server): label for label, server in self.servers.items()}
            results = {label: result for future in as_completed(futures) for label, result in [future.result()]}

        dt = datetime.now(timezone.utc).isoformat(timespec="seconds")
        for label, rows in results.items():
            if rows is None:
                continue
            server = self.servers[label]
            updated_rows = [
                {
                    "Last checked": dt,
                    "Server label": label,
                    "Server hostname": server.hostname,
                    "Server type": server.server_type,
                    **r,
                }
                for r in rows
            ]
            dbf.bulk_add_or_update_pipeline(updated_rows)
        dbf.fill_pipelines_missing_job_ids()
        dbf.update_flags()

    def load(self, pipeline_id: str, sample_id: str) -> None:
        """Load a sample on a pipeline.

        The appropriate server is found based on the pipeline, and the sample is loaded.

        Args:
            sample_id (str):
                The sample ID to load. Must exist in samples table of database
            pipeline_id (str):
                The pipeline to load the sample on. Must exist in pipelines table of database

        """
        sample = _Sample.from_id(sample_id)
        pipeline = _Pipeline.from_id(pipeline_id)
        pipeline.load(sample)

    def eject(self, pipeline_id: str, sample_id: str | None = None) -> None:
        """Eject a sample from a pipeline.

        Args:
            pipeline_id (str):
                The pipeline to eject the sample from, must exist in pipelines table of database
            sample_id (str, optional):
                Check that this sample is on the pipeline before ejecting

        """
        _Pipeline.from_id(pipeline_id).eject(sample_id)

    def submit(
        self,
        sample_id: str,
        payload: str | Path | dict,
        capacity_Ah: float | Literal["areal", "mass", "nominal"],
        comment: str = "",
    ) -> None:
        """Submit a job to a server.

        Args:
            sample_id: str
                The sample ID to submit the job for, must exist in samples table of database
            payload: str or Path or dict
                Preferably an aurora-unicycler dictionary - this is auto-converted to the right format for each cycler
                In addition, different cyclers can accept different payload formats
                (Neware) A .xml path or xml string with a Neware protocol
                (Biologic) A .mps path or mps string with a Biologic protocol
            capacity_Ah: float or str
                The capacity of the sample in Ah, if 'areal', 'mass', or 'nominal', the capacity is
                calculated from the sample information
            comment: str, optional
                A comment to add to the job in the database

        """
        sample = _Sample.from_id(sample_id)

        if isinstance(capacity_Ah, str):
            capacity_Ah = sample.get_sample_capacity(capacity_Ah)

        cycling_job = _CyclingJob(
            sample=sample,
            job_name=f"Job for sample {sample.id}",
            capacity_Ah=capacity_Ah,
            comment=comment,
        )
        cycling_job.add_payload(payload)
        cycling_job.submit()

    def cancel(self, jobid: str) -> None:
        """Cancel a job on a server.

        Args:
            jobid: str
                The job ID to cancel, must exist in jobs table of database

        Returns:
            str: The output from the server cancel command

        """
        return _CyclingJob.from_id(jobid).cancel()

    def snapshot(
        self,
        samp_or_jobid: str,
        mode: Literal["always", "new_data", "if_not_exists"] = "new_data",
    ) -> None:
        """Snapshot sample or job, download data, process, and save.

        Args:
            samp_or_jobid: str
                The sample ID or (aurora) job ID to snapshot.
            mode: str, optional
                When to make a new snapshot. Can be one of the following:
                    - 'always': Force a snapshot even if job is already done and data is downloaded.
                    - 'new_data': Snapshot if there is new data.
                    - 'if_not_exists': Snapshot only if the file doesn't exist locally.
                Default is 'new_data'.

        """
        # check if the input is a sample ID
        is_sample = dbf.is_sample(samp_or_jobid)
        if is_sample:
            jobs = [dbf.get_job_data(j) for j in dbf.get_jobs_from_sample(samp_or_jobid)]
        else:  # it's a job ID
            jobs = [dbf.get_job_data(samp_or_jobid)]
        jobs = [j for j in jobs if j is not None]
        if not jobs:
            msg = f"Sample or job ID '{samp_or_jobid}' not found in the database."
            raise ValueError(msg)

        for job in jobs:
            sample_id = job.get("Sample ID")
            job_id = job.get("Job ID")
            job_id_on_server = job.get("Job ID on server")
            if not job_id:
                continue
            if not sample_id:
                logger.warning("Job %s has no sample, skipping.", job["Job ID"])
                continue
            if not job_id_on_server:
                logger.warning("Job %s has no job ID on server, skipping.", job["Job ID"])
                continue
            # Check that sample is known
            if sample_id == "Unknown":
                logger.warning("Job %s has no sample name or payload, skipping.", job["Job ID"])
                continue

            local_save_location_processed = get_sample_folder(job["Sample ID"])

            files_exist = (local_save_location_processed / f"snapshot.{job_id}.h5").exists() or (
                local_save_location_processed / "snapshots" / f"snapshot.{job_id}.parquet"
            ).exists()
            if files_exist and mode != "always":
                if mode == "if_not_exists":
                    logger.info("Snapshot for %s already exists, skipping.", job_id)
                    continue
                if mode == "new_data" and job["Snapshot status"] is not None and job["Snapshot status"].startswith("c"):
                    logger.info("Snapshot for %s already complete, skipping.", job_id)
                    continue

            # Check that the job has started
            if job["Snapshot status"] in ["q", "qw"]:
                logger.warning("Job %s is still queued, skipping snapshot.", job_id)
                continue

            # Check that the server is accessible
            try:
                server = find_server(job["Server label"])
            except KeyError as e:
                logger.warning("Could not access server %s for job %s: %s", job["Server label"], job_id, e)
                continue

            # Snapshot the job
            try:
                new_snapshot_status = server.snapshot(sample_id, job_id, job_id_on_server)
            except FileNotFoundError as e:
                msg = (
                    f"Error snapshotting {job_id}: {e}\n"
                    "Likely the job was cancelled before starting. "
                    "Setting `Snapshot Status` to 'ce' in the database."
                )
                dbf.add_or_update_job(job_id, {"Snapshot status": "ce"})
                raise FileNotFoundError(msg) from e

            # Update the snapshot status in the database
            dt = datetime.now(timezone.utc).isoformat(timespec="seconds")
            dbf.update_results(sample_id, {"Last snapshot": dt})
            dbf.add_or_update_job(job_id, {"Last snapshot": dt, "Snapshot status": new_snapshot_status})

        # Analyse the new data (only once per sample)
        samples = [j.get("Sample ID") for j in jobs]
        unique_samples = {s for s in samples if s is not None}
        for unique_sample_id in unique_samples:
            analysis.analyse_sample(unique_sample_id)

    def _update_neware_jobids(self) -> None:
        """Update all Job IDs on Neware servers.

        Temporary measure until we have a faster way to get Job IDs from Newares
        that can run in update_db. This implementation takes ~1 second per channel.

        """
        pipelines, server_labels = dbf.get_neware_pipelines()
        for pipeline, server_label in zip(pipelines, server_labels, strict=True):
            logger.info("Updating job ID for %s on server %s", pipeline, server_label)
            server = find_server(server_label)
            assert isinstance(server, cycler_servers.NewareServer)  # noqa: S101
            jobid_on_server = server._get_job_id(pipeline)  # noqa: SLF001
            full_jobid = dbf.get_job_id_from_server(server_label, jobid_on_server)
            dbf.add_or_update_pipeline(pipeline, {"Job ID": full_jobid, "Job ID on server": jobid_on_server})