terrakio-core 0.3.0__tar.gz → 0.3.2__tar.gz

{terrakio_core-0.3.0 → terrakio_core-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: terrakio-core
-Version: 0.3.0
+Version: 0.3.2
 Summary: Core components for Terrakio API clients
 Author-email: Yupeng Chao <yupeng@haizea.com.au>
 Project-URL: Homepage, https://github.com/HaizeaAnalytics/terrakio-python-api
@@ -22,6 +22,7 @@ Requires-Dist: xarray>=2023.1.0
 Requires-Dist: shapely>=2.0.0
 Requires-Dist: geopandas>=0.13.0
 Requires-Dist: google-cloud-storage>=2.0.0
+Requires-Dist: nest_asyncio
 
 # Terrakio Core
 

{terrakio_core-0.3.0 → terrakio_core-0.3.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "terrakio-core"
-version = "0.3.0"
+version = "0.3.2"
 authors = [
     {name = "Yupeng Chao", email = "yupeng@haizea.com.au"},
 ]
@@ -29,6 +29,7 @@ dependencies = [
     "shapely>=2.0.0",
     "geopandas>=0.13.0",
     "google-cloud-storage>=2.0.0",
+    "nest_asyncio",
 ]
 
 [project.urls]

{terrakio_core-0.3.0 → terrakio_core-0.3.2}/terrakio_core/__init__.py

@@ -4,4 +4,4 @@ Terrakio Core
 Core components for Terrakio API clients.
 """
 
-__version__ = "0.3.0"
+__version__ = "0.3.2"

{terrakio_core-0.3.0 → terrakio_core-0.3.2}/terrakio_core/client.py

@@ -13,6 +13,7 @@ from shapely.geometry import shape, mapping
 from shapely.geometry.base import BaseGeometry as ShapelyGeometry
 from google.cloud import storage
 from .exceptions import APIError, ConfigurationError
+from .decorators import admin_only_params
 import logging
 import textwrap
 
@@ -129,6 +130,7 @@ class BaseClient:
             "resolution": resolution,
             **kwargs
         }
+        print("the payload is ", payload)
         request_url = f"{self.url}/geoquery"
         for attempt in range(retry + 1):
             try:
@@ -536,8 +538,8 @@ class BaseClient:
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
-    # Mass Stats methods
-    def upload_mass_stats(self, name, size, bucket, output, location=None, **kwargs):
+    @admin_only_params('location', 'force_loc', 'server')
+    def execute_job(self, name, region, output, config, overwrite=False, skip_existing=False, request_json=None, manifest_json=None, location=None, force_loc=None, server="dev-au.terrak.io"):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
             if not self.url or not self.key:
@@ -548,20 +550,8 @@ class BaseClient:
                 verify=self.verify,
                 timeout=self.timeout
             )
-        return self.mass_stats.upload_request(name, size, bucket, output, location, **kwargs)
+        return self.mass_stats.execute_job(name, region, output, config, overwrite, skip_existing, request_json, manifest_json, location, force_loc, server)
 
-    def start_mass_stats_job(self, task_id):
-        if not self.mass_stats:
-            from terrakio_core.mass_stats import MassStats
-            if not self.url or not self.key:
-                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
-            self.mass_stats = MassStats(
-                base_url=self.url,
-                api_key=self.key,
-                verify=self.verify,
-                timeout=self.timeout
-            )
-        return self.mass_stats.start_job(task_id)
 
     def get_mass_stats_task_id(self, name, stage, uid=None):
         if not self.mass_stats:
@@ -576,7 +566,7 @@ class BaseClient:
             )
         return self.mass_stats.get_task_id(name, stage, uid)
 
-    def track_mass_stats_job(self, ids=None):
+    def track_mass_stats_job(self, ids: Optional[list] = None):
         if not self.mass_stats:
             from terrakio_core.mass_stats import MassStats
             if not self.url or not self.key:
@@ -1049,6 +1039,20 @@ class BaseClient:
             )
         return self.space_management.delete_data_in_path(path, region)
     
+    def start_mass_stats_job(self, task_id):
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+        return self.mass_stats.start_job(task_id)
+
+    
     def generate_ai_dataset(
         self,
         name: str,
@@ -1118,27 +1122,58 @@ class BaseClient:
             overwrite=True
         )["task_id"]
         print("the task id is ", task_id)
+
+        # Wait for job completion
+        import time
+        
+        while True:
+            result = self.track_mass_stats_job(ids=[task_id])
+            status = result[task_id]['status']
+            print(f"Job status: {status}")
+            
+            if status == "Completed":
+                break
+            elif status == "Error":
+                raise Exception(f"Job {task_id} encountered an error")
+                
+            # Wait 30 seconds before checking again
+            time.sleep(30)
+
+        # print("the result is ", result)
+        # after all the random sample jos are done, we then start the mass stats job
         task_id = self.start_mass_stats_job(task_id)
-        print("the task id is ", task_id)
-        return task_id
+        # now we hav ethe random sampel
 
+        # print("the task id is ", task_id)
+        return task_id
 
-    def train_model(self, model_name: str, training_data: dict) -> dict:
+    def train_model(self, model_name: str, training_dataset: str, task_type: str, model_category: str, architecture: str, region: str, hyperparameters: dict = None) -> dict:
         """
         Train a model using the external model training API.
-
+        
         Args:
             model_name (str): The name of the model to train.
-            training_data (dict): Dictionary containing training data parameters.
-
+            training_dataset (str): The training dataset identifier.
+            task_type (str): The type of ML task (e.g., regression, classification).
+            model_category (str): The category of model (e.g., random_forest).
+            architecture (str): The model architecture.
+            region (str): The region identifier.
+            hyperparameters (dict, optional): Additional hyperparameters for training.
+            
         Returns:
             dict: The response from the model training API.
         """
-        endpoint = "https://modeltraining-573248941006.australia-southeast1.run.app/train_model"
         payload = {
             "model_name": model_name,
-            "training_data": training_data
+            "training_dataset": training_dataset,
+            "task_type": task_type,
+            "model_category": model_category,
+            "architecture": architecture,
+            "region": region,
+            "hyperparameters": hyperparameters
         }
+        endpoint = f"{self.url.rstrip('/')}/train_model"
+        print("the payload is ", payload)
         try:
             response = self.session.post(endpoint, json=payload, timeout=self.timeout, verify=self.verify)
             if not response.ok:
@@ -1155,35 +1190,163 @@ class BaseClient:
         except requests.RequestException as e:
             raise APIError(f"Model training request failed: {str(e)}")
 
-    def deploy_model(self, dataset: str, product:str, model_name:str, input_expression: str, model_training_job_name: str, uid: str, dates_iso8601: list):
-        # we have the dataset and we have the product, and we have the model name, we need to create a new json file and add that to the dataset as our virtual dataset
-        # upload the script to the bucket, the script should be able to download the model and do the inferencing
-        # we need to upload the the json to the to the dataset as our virtual dataset
-        # then we do nothing and wait for the user to make the request call to the explorer
-        # we should have a uniform script for the random forest deployment
-        # create a script for each model
-        # upload script to google bucket, 
-        # 
+    # Mass Stats methods
+    def combine_tiles(self, 
+        data_name: str,
+        usezarr: bool,
+        overwrite: bool,
+        output : str) -> dict:
+
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+        return self.mass_stats.combine_tiles(data_name, usezarr, overwrite, output)
+    
+
+
+    def create_dataset_file(
+        self,
+        name: str,
+        aoi: str,
+        expression: str,
+        output: str,
+        tile_size: float = 128.0,
+        crs: str = "epsg:4326",
+        res: float = 0.0001,
+        region: str = "aus",
+        to_crs: str = "epsg:4326",
+        overwrite: bool = True,
+        skip_existing: bool = False,
+        non_interactive: bool = True,
+        usezarr: bool = False,
+        poll_interval: int = 30   # seconds between job status checks
+    ) -> dict:
+        
+        from terrakio_core.generation.tiles import tiles
+        import tempfile
+        import time
+
+        body, reqs, groups = tiles(
+            name = name, 
+            aoi = aoi, 
+            expression = expression,
+            output = output,
+            tile_size = tile_size,
+            crs = crs,
+            res = res,
+            region = region,
+            to_crs = to_crs,
+            fully_cover = True,
+            overwrite = overwrite,
+            skip_existing = skip_existing,
+            non_interactive = non_interactive
+        )
+
+        # Create temp json files before upload
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tempreq:
+            tempreq.write(reqs)
+            tempreqname = tempreq.name
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tempmanifest:
+            tempmanifest.write(groups)
+            tempmanifestname = tempmanifest.name
+        
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
 
+        task_id = self.mass_stats.execute_job(
+            name=body["name"],
+            region=body["region"],
+            output=body["output"],
+            config = {},
+            overwrite=body["overwrite"],
+            skip_existing=body["skip_existing"],
+            request_json=tempreqname,
+            manifest_json=tempmanifestname,
+        )
+
+        ### Start combining tiles when generation-tiles job is done
+        start_time = time.time()
+        status = None
+        
+        while True:
+            try:
+                taskid = task_id['task_id']
+                trackinfo = self.mass_stats.track_job([taskid])
+                status = trackinfo[taskid]['status']
+                
+                # Check completion states
+                if status == 'Completed':
+                    print('Tiles generated successfully!')
+                    break
+                elif status in ['Failed', 'Cancelled', 'Error']:
+                    raise RuntimeError(f"Job {taskid} failed with status: {status}")
+                else:
+                    # Job is still running
+                    elapsed_time = time.time() - start_time
+                    print(f"Job status: {status} - Elapsed time: {elapsed_time:.1f}s", end='\r')
+                    
+                    # Sleep before next check
+                    time.sleep(poll_interval)
+                    
+                    
+            except KeyboardInterrupt:
+                print(f"\nInterrupted! Job {taskid} is still running in the background.")
+                raise
+            except Exception as e:
+                print(f"\nError tracking job: {e}")
+                raise
+
+        # Clean up temporary files
+        import os
+        os.unlink(tempreqname)
+        os.unlink(tempmanifestname)
+
+
+        # Start combining tiles
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+        
+        return self.mass_stats.combine_tiles(body["name"], usezarr, body["overwrite"], body["output"])
+
+    def deploy_model(self, dataset: str, product:str, model_name:str, input_expression: str, model_training_job_name: str, uid: str, dates_iso8601: list):
         script_content = self._generate_script(model_name, product, model_training_job_name, uid)
-        # self.create_dataset(collection = "terrakio-datasets", input = input, )
-        # we have the script, we need to upload it to the bucket
         script_name = f"{product}.py"
-        print("the script content is ", script_content)
-        print("the script name is ", script_name)
         self._upload_script_to_bucket(script_content, script_name, model_training_job_name, uid)
-        # after uploading the script, we need to create a new virtual dataset
         self._create_dataset(name = dataset, collection = "terrakio-datasets", products = [product], path = f"gs://terrakio-mass-requests/{uid}/{model_training_job_name}/inference_scripts", input = input_expression, dates_iso8601 = dates_iso8601, padding = 0)
 
     def _generate_script(self, model_name: str, product: str, model_training_job_name: str, uid: str) -> str:
         return textwrap.dedent(f'''
             import logging
             from io import BytesIO
-            from google.cloud import storage
-            from onnxruntime import InferenceSession
+
             import numpy as np
+            import pandas as pd
             import xarray as xr
-            import datetime
+            from google.cloud import storage
+            from onnxruntime import InferenceSession
 
             logging.basicConfig(
                 level=logging.INFO
@@ -1191,54 +1354,122 @@ class BaseClient:
 
             def get_model():
                 logging.info("Loading model for {model_name}...")
-                
+
                 client = storage.Client()
                 bucket = client.get_bucket('terrakio-mass-requests')
                 blob = bucket.blob('{uid}/{model_training_job_name}/models/{model_name}.onnx')
-                
+
                 model = BytesIO()
                 blob.download_to_file(model)
                 model.seek(0)
-                
+
                 session = InferenceSession(model.read(), providers=["CPUExecutionProvider"])
                 return session
 
             def {product}(*bands, model):
                 logging.info("start preparing data")
+                print("the bands are ", bands)
                 
-                original_shape = bands[0].shape
-                logging.info(f"Original shape: {{original_shape}}")
+                data_arrays = list(bands)
                 
-                transformed_bands = []
-                for band in bands:
-                    transformed_band = band.values.reshape(-1,1)
-                    transformed_bands.append(transformed_band)
+                print("the data arrays are ", [da.name for da in data_arrays])
                 
-                input_data = np.hstack(transformed_bands)
+                reference_array = data_arrays[0]
+                original_shape = reference_array.shape
+                logging.info(f"Original shape: {{original_shape}}")
                 
+                if 'time' in reference_array.dims:
+                    time_coords = reference_array.coords['time']
+                    if len(time_coords) == 1:
+                        output_timestamp = time_coords[0]
+                    else:
+                        years = [pd.to_datetime(t).year for t in time_coords.values]
+                        unique_years = set(years)
+                        
+                        if len(unique_years) == 1:
+                            year = list(unique_years)[0]
+                            output_timestamp = pd.Timestamp(f"{{year}}-01-01")
+                        else:
+                            latest_year = max(unique_years)
+                            output_timestamp = pd.Timestamp(f"{{latest_year}}-01-01")
+                else:
+                    output_timestamp = pd.Timestamp("1970-01-01")
+
+                averaged_bands = []
+                for data_array in data_arrays:
+                    if 'time' in data_array.dims:
+                        averaged_band = np.mean(data_array.values, axis=0)
+                        logging.info(f"Averaged band from {{data_array.shape}} to {{averaged_band.shape}}")
+                    else:
+                        averaged_band = data_array.values
+                        logging.info(f"No time dimension, shape: {{averaged_band.shape}}")
+
+                    flattened_band = averaged_band.reshape(-1, 1)
+                    averaged_bands.append(flattened_band)
+
+                input_data = np.hstack(averaged_bands)
+
                 logging.info(f"Final input shape: {{input_data.shape}}")
-                
+
                 output = model.run(None, {{"float_input": input_data.astype(np.float32)}})[0]
-                
+
                 logging.info(f"Model output shape: {{output.shape}}")
 
-                output_reshaped = output.reshape(original_shape)
+                if len(original_shape) >= 3:
+                    spatial_shape = original_shape[1:]
+                else:
+                    spatial_shape = original_shape
+
+                output_reshaped = output.reshape(spatial_shape)
+
+                output_with_time = np.expand_dims(output_reshaped, axis=0)
+
+                if 'time' in reference_array.dims:
+                    spatial_dims = [dim for dim in reference_array.dims if dim != 'time']
+                    spatial_coords = {{dim: reference_array.coords[dim] for dim in spatial_dims if dim in reference_array.coords}}
+                else:
+                    spatial_dims = list(reference_array.dims)
+                    spatial_coords = dict(reference_array.coords)
+
                 result = xr.DataArray(
-                    data=output_reshaped,
-                    dims=bands[0].dims,
-                    coords=bands[0].coords
+                    data=output_with_time.astype(np.float32),
+                    dims=['time'] + list(spatial_dims),
+                    coords={
+                        'time': [output_timestamp.values],
+                        'y': spatial_coords['y'].values,
+                        'x': spatial_coords['x'].values
+                    }
                 )
-                
                 return result
             ''').strip()
-
+    
     def _upload_script_to_bucket(self, script_content: str, script_name: str, model_training_job_name: str, uid: str):
         """Upload the generated script to Google Cloud Storage"""
 
         client = storage.Client()
         bucket = client.get_bucket('terrakio-mass-requests')
         blob = bucket.blob(f'{uid}/{model_training_job_name}/inference_scripts/{script_name}')
-        # the first layer is the uid, the second layer is the model training job name
         blob.upload_from_string(script_content, content_type='text/plain')
         logging.info(f"Script uploaded successfully to {uid}/{model_training_job_name}/inference_scripts/{script_name}")
 
+
+
+
+    def download_file_to_path(self, job_name, stage, file_name, output_path):
+        if not self.mass_stats:
+            from terrakio_core.mass_stats import MassStats
+            if not self.url or not self.key:
+                raise ConfigurationError("Mass Stats client not initialized. Make sure API URL and key are set.")
+            self.mass_stats = MassStats(
+                base_url=self.url,
+                api_key=self.key,
+                verify=self.verify,
+                timeout=self.timeout
+            )
+
+        # fetch bucket info based on job name and stage
+
+        taskid = self.mass_stats.get_task_id(job_name, stage).get('task_id')
+        trackinfo = self.mass_stats.track_job([taskid])
+        bucket = trackinfo[taskid]['bucket']
+        return self.mass_stats.download_file(job_name, bucket, file_name, output_path)

{terrakio_core-0.3.0 → terrakio_core-0.3.2}/terrakio_core/dataset_management.py

@@ -83,10 +83,63 @@ class DatasetManagement:
         except requests.RequestException as e:
             raise APIError(f"Request failed: {str(e)}")
 
+    # def create_dataset(self, name: str, collection: str = "terrakio-datasets", **kwargs) -> Dict[str, Any]:
+    #     """
+    #     Create a new dataset.
+        
+    #     Args:
+    #         name: Name of the dataset (required)
+    #         collection: Dataset collection (default: 'terrakio-datasets')
+    #         **kwargs: Additional dataset parameters including:
+    #             - products: List of products
+    #             - dates_iso8601: List of dates
+    #             - bucket: Storage bucket
+    #             - path: Storage path
+    #             - data_type: Data type
+    #             - no_data: No data value
+    #             - l_max: Maximum level
+    #             - y_size: Y size
+    #             - x_size: X size
+    #             - proj4: Projection string
+    #             - abstract: Dataset abstract
+    #             - geotransform: Geotransform parameters
+                
+    #     Returns:
+    #         Created dataset information
+            
+    #     Raises:
+    #         APIError: If the API request fails
+    #     """
+    #     endpoint = f"{self.api_url}/datasets"
+    #     params = {"collection": collection}
+    #     # Create payload with required name parameter
+    #     payload = {"name": name}
+        
+    #     # Add optional parameters if provided
+    #     for param in ["products", "dates_iso8601", "bucket", "path", "data_type", 
+    #                  "no_data", "l_max", "y_size", "x_size", "proj4", "abstract", "geotransform", "input"]:
+    #         if param in kwargs:
+    #             payload[param] = kwargs[param]
+        
+    #     try:
+    #         response = self.session.post(
+    #             endpoint,
+    #             params=params,
+    #             json=payload,
+    #             timeout=self.timeout,
+    #             verify=self.verify
+    #         )
+            
+    #         if not response.ok:
+    #             raise APIError(f"API request failed: {response.status_code} {response.reason}")
+    #         return response.json()
+    #     except requests.RequestException as e:
+    #         raise APIError(f"Request failed: {str(e)}")
+
     def create_dataset(self, name: str, collection: str = "terrakio-datasets", **kwargs) -> Dict[str, Any]:
         """
         Create a new dataset.
-        
+                    
         Args:
             name: Name of the dataset (required)
             collection: Dataset collection (default: 'terrakio-datasets')
@@ -103,24 +156,23 @@ class DatasetManagement:
                 - proj4: Projection string
                 - abstract: Dataset abstract
                 - geotransform: Geotransform parameters
-                
+                - padding: Padding value
+                            
         Returns:
             Created dataset information
-            
+                        
         Raises:
             APIError: If the API request fails
         """
         endpoint = f"{self.api_url}/datasets"
         params = {"collection": collection}
-        # Create payload with required name parameter
         payload = {"name": name}
-        
-        # Add optional parameters if provided
-        for param in ["products", "dates_iso8601", "bucket", "path", "data_type", 
-                     "no_data", "l_max", "y_size", "x_size", "proj4", "abstract", "geotransform", "input"]:
+                    
+        for param in ["products", "dates_iso8601", "bucket", "path", "data_type",
+                        "no_data", "l_max", "y_size", "x_size", "proj4", "abstract", "geotransform", "input", "padding"]:
             if param in kwargs:
                 payload[param] = kwargs[param]
-        
+                    
         try:
             response = self.session.post(
                 endpoint,
@@ -129,7 +181,7 @@ class DatasetManagement:
                 timeout=self.timeout,
                 verify=self.verify
             )
-            
+                            
             if not response.ok:
                 raise APIError(f"API request failed: {response.status_code} {response.reason}")
             return response.json()

terrakio_core-0.3.2/terrakio_core/decorators.py

@@ -0,0 +1,18 @@
+# terrakio_core/decorators.py
+def admin_only_params(*restricted_params):
+    """
+    Decorator factory for restricting method parameters to admin users only.
+    """
+    def decorator(func):
+        def wrapper(self, *args, **kwargs):
+            if hasattr(self, '_is_admin') and self._is_admin:
+                return func(self, *args, **kwargs)
+            
+            admin_params_used = set(kwargs.keys()) & set(restricted_params)
+            if admin_params_used:
+                raise PermissionError(f"Parameters {admin_params_used} are only available to admin users")
+            
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k not in restricted_params}
+            return func(self, *args, **filtered_kwargs)
+        return wrapper
+    return decorator