Repository: oumi-ai/oumi
File: src/oumi/launcher/clients/sky_client.py (Python)

Code before:
    Fetches the override value from the OUMI_USE_SPOT_VM environment variable
    if specified.

    Returns:
        The override value if specified, or `None`.
    """
    _ENV_VAR_NAME = "OUMI_USE_SPOT_VM"
    s = os.environ.get(_ENV_VAR_NAME, "")
    mode = s.lower().replace("-", "").replace("_", "").strip()
    if not mode or mode in ("config",):
        return None
    bool_result = try_str_to_bool(mode)
    if bool_result is not None:
        return bool_result
    if mode in ("spot", "preemptible", "preemptable"):
        return True
    elif mode in ("nonspot", "nonpreemptible", "nonpreemptable"):
        return False
    raise ValueError(f"{_ENV_VAR_NAME} has unsupported value: '{s}'.")


def _convert_job_to_task(job: JobConfig) -> "sky.Task":
    """Converts a JobConfig to a sky.Task."""
    # Delay sky import: https://github.com/oumi-ai/oumi/issues/1605
    import sky

    sky_cloud = _get_sky_cloud_from_job(job)
    use_spot_vm = _get_use_spot_vm_override()
    if use_spot_vm is None:
        use_spot_vm = job.resources.use_spot
    elif use_spot_vm != job.resources.use_spot:
        logger.info(f"Set use_spot={use_spot_vm} based on 'OUMI_USE_SPOT_VM' override.")

    resources = sky.Resources(
        cloud=sky_cloud,
        instance_type=job.resources.instance_type,
        cpus=job.resources.cpus,
        memory=job.resources.memory,
        accelerators=job.resources.accelerators,
        use_spot=use_spot_vm,
        region=job.resources.region,
        zone=job.resources.zone,
        disk_size=job.resources.disk_size,
        disk_tier=job.resources.disk_tier,
    )
    sky_task = sky.Task(
        name=job.name,
        setup=job.setup,
        run=job.run,
        envs=job.envs,
        workdir=job.working_dir,

Review comment (suggestion):
Your other files have this comment for all sky imports. Do you want to do the same in this file? If you're only going to have it in one place, consider putting it at the top or in the SkyClient class for discoverability.

Code after:
    """Determines whether to override `use_spot_vm` setting based on OUMI_USE_SPOT_VM.

    Fetches the override value from the OUMI_USE_SPOT_VM environment variable
    if specified.

    Returns:
        The override value if specified, or `None`.
    """
    _ENV_VAR_NAME = "OUMI_USE_SPOT_VM"
    s = os.environ.get(_ENV_VAR_NAME, "")
    mode = s.lower().replace("-", "").replace("_", "").strip()
    if not mode or mode in ("config",):
        return None
    bool_result = try_str_to_bool(mode)
    if bool_result is not None:
        return bool_result
    if mode in ("spot", "preemptible", "preemptable"):
        return True
    elif mode in ("nonspot", "nonpreemptible", "nonpreemptable"):
        return False
    raise ValueError(f"{_ENV_VAR_NAME} has unsupported value: '{s}'.")


def _convert_job_to_task(job: JobConfig) -> "sky.Task":
    """Converts a JobConfig to a sky.Task."""
    # Delay sky import: https://github.com/oumi-ai/oumi/issues/1605
    import sky

    sky_cloud = _get_sky_cloud_from_job(job)
    use_spot_vm = _get_use_spot_vm_override()
    if use_spot_vm is None:
        use_spot_vm = job.resources.use_spot
    elif use_spot_vm != job.resources.use_spot:
        logger.info(f"Set use_spot={use_spot_vm} based on 'OUMI_USE_SPOT_VM' override.")

    resources = sky.Resources(
        cloud=sky_cloud,
        instance_type=job.resources.instance_type,
        cpus=job.resources.cpus,
        memory=job.resources.memory,
        accelerators=job.resources.accelerators,
        use_spot=use_spot_vm,
        region=job.resources.region,
        zone=job.resources.zone,
        disk_size=job.resources.disk_size,
        disk_tier=job.resources.disk_tier,
    )
    sky_task = sky.Task(
        name=job.name,
        setup=job.setup,
        run=job.run,