here is the full code block context for you
# Import necessary libraries
import numpy as np
import rioxarray
import xarray as xr
import copernicusmarine
import datetime
import s3fs
import os
import coiled
from prefect import task, flow
import pandas as pd
import rasterio
import io
import ftplib
import tempfile
import rasterio
from rasterio.fill import fillnodata
from zoneinfo import ZoneInfo
from prefect.blocks.system import Secret
# Initialize S3 filesystem
fs = s3fs.S3FileSystem(anon=False)
# ------------------------- JAXA FLOW (Coiled) -------------------------
@task
@coiled.function(
name="prefect_jaxa_hourly_sst_median",
container="bathymaps/geodocker:v1.4",
region="ap-southeast-2",
vm_type="m7a.4xlarge",
arm=False,
n_workers=1,
idle_timeout="10m",
environ={"CPL_VSIL_USE_TEMP_FILE_FOR_RANDOM_WRITE": "YES"},
)
def run_jaxa_flow():
def _ftp_connect(host, username, password):
ftp = ftplib.FTP(host, timeout=60)
ftp.login(username, password)
return ftp
def _ftp_disconnect(ftp):
try:
ftp.quit()
except ftplib.all_errors:
pass
def _list_latest_hour_files(ftp, base_path, keep=6):
# latest month
month_folders = sorted(
[p for p in ftp.nlst(base_path) if p.strip().split('/')[-1].isdigit()],
reverse=True
)
if not month_folders:
return None, None, []
latest_month_path = month_folders[0]
latest_month_leaf = latest_month_path.strip('/').split('/')[-1]
# latest day
ftp.cwd(latest_month_path)
day_folders = sorted(
[d for d in ftp.nlst() if d.strip().split('/')[-1].isdigit()],
reverse=True
)
if not day_folders:
return latest_month_leaf, None, []
latest_day_path = day_folders[0]
latest_day_leaf = latest_day_path.strip('/').split('/')[-1]
# files (build absolute remote paths)
ftp.cwd(latest_day_path)
file_names = [f for f in ftp.nlst() if f.endswith(".nc") and not f.endswith(".md5")]
files_full = [
f"{base_path.rstrip('/')}/{latest_month_leaf}/{latest_day_leaf}/{fn.split('/')[-1]}"
for fn in sorted(file_names, reverse=True)[:keep]
]
return latest_month_leaf, latest_day_leaf, files_full
def list_latest_files(ftp_host: str, ftp_user: str, ftp_pass: str, base_path: str, keep: int = 6):
ftp = _ftp_connect(ftp_host, ftp_user, ftp_pass)
try:
latest_month, latest_day, files = _list_latest_hour_files(ftp, base_path, keep=keep)
print(f"Latest: month={latest_month}, day={latest_day}, n_files={len(files)}")
return latest_month, latest_day, files
finally:
_ftp_disconnect(ftp)
def fetch_dataset(ftp_host: str, ftp_user: str, ftp_pass: str, base_path: str, remote_relpath: str):
ftp = _ftp_connect(ftp_host, ftp_user, ftp_pass)
try:
buf = io.BytesIO()
# Absolute path from root: /<base_path>/<month>/<day>/<file>
ftp.retrbinary(f"RETR /{base_path.rstrip('/')}/{remote_relpath}", buf.write)
buf.seek(0)
ds = xr.open_dataset(
buf,
decode_times=True,
decode_timedelta=True
) # let xarray auto-detect engine
print(f"Loaded dataset from {remote_relpath} with vars: {list(ds.data_vars)}")
return ds
finally:
_ftp_disconnect(ftp)
def process_sst_data(ds_list):
#logger = get_run_logger()
if not ds_list:
print("No datasets to process")
return None
print("Processing SST data now...")
# Concatenate and process
out = xr.concat(ds_list, dim="time")
sst = (
out["sea_surface_temperature"]
.sortby("lon", ascending=True)
.sel(lon=slice(107, 160), lat=slice(-8, -48))
- 273.15
)
median_sst = sst.median(dim="time", skipna=True)
median_sst = median_sst.sortby("lat", ascending=False)
median_sst = median_sst.rio.set_spatial_dims(x_dim="lon", y_dim="lat", inplace=False)
median_sst = median_sst.rio.write_crs("EPSG:4326", inplace=False)
median_sst = median_sst.rio.write_nodata(np.nan, encoded=True, inplace=False)
print(f"Processed SST data: shape={median_sst.shape}")
return median_sst
def export_to_geotiff(
data,
file_type: str,
timestamp: str,
s3_prefix: str = "jaxa_sst/hourly",
object_name: str | None = None,
):
print(f"Exporting data to GeoTIFF: type={file_type}, timestamp={timestamp}")
s3 = s3fs.S3FileSystem(anon=False)
if data is None:
print("No data to export")
return None
data = data.astype("float32")
data = data.rio.write_crs("EPSG:4326", inplace=False)
data = data.sortby("lat", ascending=False)
data = data.rio.write_nodata(np.nan, encoded=True, inplace=False)
data = data.rio.set_spatial_dims(x_dim="lon", y_dim="lat", inplace=False)
arr = data.values.astype("float32")
transform = data.rio.transform()
profile = {
"driver": "GTiff",
"dtype": "float32",
"count": 1,
"width": data.rio.width,
"height": data.rio.height,
"crs": "EPSG:4326",
"transform": transform,
"nodata": np.nan,
"compress": "deflate",
"predictor": 2,
"tiled": True,
"blockxsize": 512,
"blockysize": 512,
}
key = object_name or f"sst-himawari-nrt-{file_type}-{timestamp}.tif"
s3_key = f"{s3_prefix.rstrip('/')}/{key}"
bucket = "copernicusbathymaps"
path = f"s3://{bucket}/{s3_key}"
with rasterio.open(path, "w", **profile) as dst:
dst.write(arr, 1)
dst.set_band_description(1, "median_sst_hourly_celcius")
# load secrets from Prefect
ftp_host = '
ftp.ptree.jaxa.jp'
ftp_user = Secret.load("jaxa-user").get()
ftp_pass = Secret.load("jaxa-pw").get()
base_path = 'pub/himawari/L2/SST/v201_nc4_normal_nrt'
# Step 1: List latest files
print("Listing latest files...")
latest_month, latest_day, files = list_latest_files(
ftp_host, ftp_user, ftp_pass, base_path, keep=6
)
if not files:
print("No files found; exiting flow.")
return
# Step 2: Fetch datasets
datasets = []
times = []
for fullpath in files:
# fullpath is "<base_path>/<month>/<day>/<filename>"
rel = "/".join(fullpath.split("/")[len(base_path.strip('/').split('/')):]) # month/day/filename
ds = fetch_dataset(ftp_host, ftp_user, ftp_pass, base_path, rel)
ts = ds.time[0].values
times.append(
pd.to_datetime(ts).to_pydatetime())
datasets.append(ds)
# Step 3: Process SST data
median_sst = process_sst_data(datasets)
# Step 4: Export to S3
# run related timestamp
run_timestamp1 = datetime.datetime.now() + datetime.timedelta(hours=10)
run_timestamp2 = run_timestamp1.strftime("%Y%m%d-%H%M")
# file related timestamp (from latest file)
pyts = times[0]
# Convert file timestamp to UTC+10
pyts = (
pyts.replace(tzinfo=datetime.timezone.utc)
.astimezone(ZoneInfo("Australia/Brisbane"))
)
file_timestamp = pyts.strftime("%Y%m%d-%H%M")
s3_key = export_to_geotiff(median_sst, "hourly", file_timestamp)
print(f"Flow completed successfully. Output: {s3_key}")
return s3_key
#---- Flow to orchestrate the tasks ----#
@flow(name="JAXA Flow", log_prints=True)
def jaxa_flow():
# Run JAXA once (via Coiled) after Copernicus hours
jaxa_future = run_jaxa_flow.submit()
try:
jaxa_result = jaxa_future.result()
print(f"✅ JAXA wrote
s3://copernicusbathymaps/{jaxa_result}")
except Exception as e:
print(f"⚠️ JAXA failed: {e}")
if
name == "__main__":
jaxa_flow.deploy(
name="jaxa_flow",
work_pool_name="bathymaps-coiled-pool",
image="bathymaps/geodocker:v1.4",
build=False,
cron="0 * * * *",
job_variables={
"arm":False
}
)