starflow-py 0.49.0__tar.gz → 0.56.0__tar.gz

{starflow_py-0.49.0 → starflow_py-0.56.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: starflow-py
-Version: 0.49.0
+Version: 0.56.0
 Summary: Python package for starflow.
 Author-email: Alexandre Lazar <alexandre.lazar@gmail.com>
 License: 

starflow_py-0.56.0/src/starflow/VERSION

@@ -0,0 +1 @@
+0.56.0

{starflow_py-0.49.0 → starflow_py-0.56.0}/src/starflow/custom_operators/k8s_branch_operator.py

@@ -38,6 +38,18 @@ class starflowKubernetesBranchPodOperator(KubernetesPodOperator, BaseBranchOpera
         container_resources: Optional[Dict] = None,
         **k8s_operator_kwargs
     ):
+        import logging
+        logger = logging.getLogger(__name__)
+        
+        logger.info("=" * 80)
+        logger.info(f"[StarflowK8sBranchOperator] INIT METHOD CALLED - task_id: {task_id}")
+        logger.info(f"[StarflowK8sBranchOperator] dag_id: {dag_id}")
+        logger.info(f"[StarflowK8sBranchOperator] piece_name: {piece_name}")
+        logger.info(f"[StarflowK8sBranchOperator] repository_url: {repository_url}")
+        logger.info(f"[StarflowK8sBranchOperator] repository_version: {repository_version}")
+        logger.info(f"[StarflowK8sBranchOperator] workspace_id: {workspace_id}")
+        logger.info("=" * 80)
+        
         self.logger = get_configured_logger("starflowKubernetesBranchPodOperator")
         self.task_id = task_id
         self.piece_name = piece_name
@@ -214,9 +226,27 @@ class starflowKubernetesBranchPodOperator(KubernetesPodOperator, BaseBranchOpera
         self.task_id_replaced = self.task_id.lower().replace("_", "-") # doing this because airflow doesn't allow underscores and upper case in mount names and max len is 63
         self.shared_storage_base_mount_path = '/home/shared_storage'
 
+        # Fix pod name to comply with Kubernetes naming rules (no underscores)
+        if pod.metadata.name:
+            # Replace underscores with hyphens and ensure it's lowercase
+            pod.metadata.name = pod.metadata.name.lower().replace("_", "-")
+            # Ensure it starts and ends with alphanumeric characters
+            pod.metadata.name = pod.metadata.name.strip("-")
+            if not pod.metadata.name[0].isalnum():
+                pod.metadata.name = "pod-" + pod.metadata.name
+            if not pod.metadata.name[-1].isalnum():
+                pod.metadata.name = pod.metadata.name + "-pod"
+
         pod.metadata.namespace = "airflow"
         pod.spec.image_pull_secrets = [k8s.V1LocalObjectReference('ghcr-login-secret')]
 
+        # Override xcom sidecar image to use Harbor registry instead of quay.io
+        if pod.spec.containers:
+            for container in pod.spec.containers:
+                if container.name == "airflow-xcom-sidecar":
+                    container.image = "harbor.atlassistant.com/public/alpine:latest"
+                    self.logger.info(f"[StarflowK8sBranchOperator] Updated xcom sidecar image to: {container.image}")
+
         # if not self.workflow_shared_storage or self.workflow_shared_storage.mode.name == 'none':
         #     return pod
         # if self.workflow_shared_storage.source.name in ["aws_s3", "gcs"]:
@@ -226,7 +256,6 @@ class starflowKubernetesBranchPodOperator(KubernetesPodOperator, BaseBranchOpera
 
         return pod
 
-
     def add_local_shared_storage_volumes(self, pod: k8s.V1Pod) -> k8s.V1Pod:
         """Adds local shared storage volumes to the pod."""
         pod_cp = copy.deepcopy(pod)
@@ -397,6 +426,11 @@ class starflowKubernetesBranchPodOperator(KubernetesPodOperator, BaseBranchOpera
         piece_repository_data = self.starflow_client.get_piece_repositories_from_workspace_id(
             params=params
         ).json()
+        
+        # Check if any piece repositories were found
+        if not piece_repository_data.get("data") or len(piece_repository_data["data"]) == 0:
+            return {}
+        
         secrets_response = self.starflow_client.get_piece_secrets(
             piece_repository_id=piece_repository_data["data"][0]["id"],
             piece_name=piece_name
@@ -467,8 +501,26 @@ class starflowKubernetesBranchPodOperator(KubernetesPodOperator, BaseBranchOpera
         Runs at the begining of the execute method.
         Pass extra arguments and configuration as environment variables to the pod
         """
+        import logging
+        logger = logging.getLogger(__name__)
+        
         # Fetch upstream tasks ids and save them in an ENV var
-        upstream_task_ids = [t.task_id for t in self.get_direct_relatives(upstream=True)]
+        # Check if operator is properly assigned to a DAG
+        try:
+            # Try to access the dag property - this will raise RuntimeError if not assigned
+            dag = self.dag
+            if dag is None:
+                logger.warning("[StarflowK8sBranchOperator] Operator DAG is None, skipping upstream task IDs")
+                upstream_task_ids = []
+            else:
+                upstream_task_ids = [t.task_id for t in self.get_direct_relatives(upstream=True)]
+                logger.info(f"[StarflowK8sBranchOperator] Upstream task IDs: {upstream_task_ids}")
+        except RuntimeError as e:
+            if "has not been assigned to a DAG yet" in str(e):
+                logger.warning("[StarflowK8sBranchOperator] Operator not assigned to a DAG, skipping upstream task IDs")
+                upstream_task_ids = []
+            else:
+                raise
         self.env_vars.append({
             'name': 'AIRFLOW_UPSTREAM_TASKS_IDS',
             'value': str(upstream_task_ids),
@@ -518,19 +570,49 @@ class starflowKubernetesBranchPodOperator(KubernetesPodOperator, BaseBranchOpera
         """
         Code from here onward is executed by the Worker and not by the Scheduler.
         """
-        # TODO change url based on platform configuration
-        self.starflow_client = starflowBackendRestClient(base_url="http://starflow-sa-rest-service.prochain-starflow:8000/")
-        self._prepare_execute_environment(context=context)
-        remote_pod = None
+        import logging
+        logger = logging.getLogger(__name__)
+        
+        # Log at the very beginning to confirm execute method is called
+        logger.info("=" * 80)
+        logger.info(f"[StarflowK8sBranchOperator] EXECUTE METHOD CALLED - task_id: {self.task_id}")
+        logger.info(f"[StarflowK8sBranchOperator] Context type: {type(context)}")
+        logger.info(f"[StarflowK8sBranchOperator] Context keys: {list(context.keys()) if hasattr(context, 'keys') else 'No keys method'}")
+        logger.info("=" * 80)
+        
         try:
+            logger.info(f"[StarflowK8sBranchOperator] Starting execution for task_id: {self.task_id}")
+            
+            # Initialize remote_pod early to avoid UnboundLocalError in finally block
+            remote_pod = None
+            
+            # TODO change url based on platform configuration
+            logger.info("[StarflowK8sBranchOperator] Initializing Starflow backend client...")
+            self.starflow_client = starflowBackendRestClient(base_url="http://starflow-rest-service.starflow:8000/")
+            logger.info("[StarflowK8sBranchOperator] Starflow backend client initialized successfully")
+            
+            logger.info("[StarflowK8sBranchOperator] Preparing execute environment...")
+            self._prepare_execute_environment(context=context)
+            logger.info("[StarflowK8sBranchOperator] Execute environment prepared successfully")
+            logger.info("[StarflowK8sBranchOperator] Building pod request object...")
             self.pod_request_obj = self.build_pod_request_obj(context)
+            logger.info("[StarflowK8sBranchOperator] Pod request object built successfully")
+            
+            logger.info("[StarflowK8sBranchOperator] Creating/getting pod...")
             self.pod = self.get_or_create_pod(  # must set `self.pod` for `on_kill`
                 pod_request_obj=self.pod_request_obj,
                 context=context,
             )
+            logger.info(f"[StarflowK8sBranchOperator] Pod created successfully: {self.pod.metadata.name}")
+            
             # get remote pod for use in cleanup methods
+            logger.info("[StarflowK8sBranchOperator] Finding remote pod...")
             remote_pod = self.find_pod(self.pod.metadata.namespace, context=context)
+            logger.info("[StarflowK8sBranchOperator] Remote pod found successfully")
+            
+            logger.info("[StarflowK8sBranchOperator] Awaiting pod start...")
             self.await_pod_start(pod=self.pod)
+            logger.info("[StarflowK8sBranchOperator] Pod started successfully")
 
             if self.get_logs:
                 self.pod_manager.fetch_container_logs(

{starflow_py-0.49.0 → starflow_py-0.56.0}/src/starflow/custom_operators/k8s_operator.py

@@ -34,6 +34,18 @@ class starflowKubernetesPodOperator(KubernetesPodOperator):
         container_resources: Optional[Dict] = None,
         **k8s_operator_kwargs
     ):
+        import logging
+        logger = logging.getLogger(__name__)
+        
+        logger.info("=" * 80)
+        logger.info(f"[StarflowK8sOperator] INIT METHOD CALLED - task_id: {task_id}")
+        logger.info(f"[StarflowK8sOperator] dag_id: {dag_id}")
+        logger.info(f"[StarflowK8sOperator] piece_name: {piece_name}")
+        logger.info(f"[StarflowK8sOperator] repository_url: {repository_url}")
+        logger.info(f"[StarflowK8sOperator] repository_version: {repository_version}")
+        logger.info(f"[StarflowK8sOperator] workspace_id: {workspace_id}")
+        logger.info("=" * 80)
+        
         self.logger = get_configured_logger("starflowKubernetesPodOperator")
         self.task_id = task_id
         self.piece_name = piece_name
@@ -45,6 +57,10 @@ class starflowKubernetesPodOperator(KubernetesPodOperator):
         self.workflow_shared_storage = workflow_shared_storage
         self.piece_source_image = k8s_operator_kwargs["image"]
         self.username_var = username
+        
+        # Initialize attributes needed by build_pod_request_obj
+        self.task_id_replaced = self.task_id.lower().replace("_", "-") # doing this because airflow doesn't allow underscores and upper case in mount names and max len is 63
+        self.shared_storage_base_mount_path = '/home/shared_storage'
 
         # Environment variables
         pod_env_vars = {
@@ -195,15 +211,31 @@ class starflowKubernetesPodOperator(KubernetesPodOperator):
     def build_pod_request_obj(self, context: Optional['Context'] = None) -> k8s.V1Pod:
         """
         We override this method to add the shared storage to the pod.
-        This function runs after our own self.execute, by super().execute()
+        This function is called by the parent class's execute method.
         """
         pod = super().build_pod_request_obj(context)
-        self.task_id_replaced = self.task_id.lower().replace("_", "-") # doing this because airflow doesn't allow underscores and upper case in mount names and max len is 63
-        self.shared_storage_base_mount_path = '/home/shared_storage'
+
+        # Fix pod name to comply with Kubernetes naming rules (no underscores)
+        if pod.metadata.name:
+            # Replace underscores with hyphens and ensure it's lowercase
+            pod.metadata.name = pod.metadata.name.lower().replace("_", "-")
+            # Ensure it starts and ends with alphanumeric characters
+            pod.metadata.name = pod.metadata.name.strip("-")
+            if not pod.metadata.name[0].isalnum():
+                pod.metadata.name = "pod-" + pod.metadata.name
+            if not pod.metadata.name[-1].isalnum():
+                pod.metadata.name = pod.metadata.name + "-pod"
 
         pod.metadata.namespace = "airflow"
         pod.spec.image_pull_secrets = [k8s.V1LocalObjectReference('ghcr-login-secret')]
 
+        # Override xcom sidecar image to use Harbor public registry instead of quay.io
+        if pod.spec.containers:
+            for container in pod.spec.containers:
+                if container.name == "airflow-xcom-sidecar":
+                    container.image = "harbor.atlassistant.com/public/alpine:latest"
+                    self.logger.info(f"[StarflowK8sOperator] Updated xcom sidecar image to: {container.image}")
+
         # if not self.workflow_shared_storage or self.workflow_shared_storage.mode.name == 'none':
         #     return pod
         # if self.workflow_shared_storage.source.name in ["aws_s3", "gcs"]:
@@ -213,7 +245,6 @@ class starflowKubernetesPodOperator(KubernetesPodOperator):
 
         return pod
 
-
     def add_local_shared_storage_volumes(self, pod: k8s.V1Pod) -> k8s.V1Pod:
         """Adds local shared storage volumes to the pod."""
         pod_cp = copy.deepcopy(pod)
@@ -373,30 +404,74 @@ class starflowKubernetesPodOperator(KubernetesPodOperator):
         source: str = 'github'
     ) -> Dict[str, Any]:
         """Get piece secrets values from starflow API"""
-        params = {
-            "workspace_id": self.workspace_id,
-            "url": repository_url,
-            "version": repository_version,
-            'source': source,
-            "page": 0,
-            "page_size": 1,
-        }
-        piece_repository_data = self.starflow_client.get_piece_repositories_from_workspace_id(
-            params=params
-        ).json()
-        secrets_response = self.starflow_client.get_piece_secrets(
-            piece_repository_id=piece_repository_data["data"][0]["id"],
-            piece_name=piece_name
-        )
-        if secrets_response.status_code != 200:
-            raise Exception(f"Error getting piece secrets: {secrets_response.json()}")
-
-        piece_secrets = {}
-        for e in secrets_response.json():
-            if not e.get('value') and not e.get('required'):
-                continue
-            piece_secrets[e.get('name')] = e.get('value')
-        return piece_secrets
+        import logging
+        logger = logging.getLogger(__name__)
+        
+        try:
+            logger.info(f"[StarflowK8sOperator] Getting piece secrets for piece: {piece_name}")
+            logger.info(f"[StarflowK8sOperator] Repository URL: {repository_url}")
+            logger.info(f"[StarflowK8sOperator] Repository version: {repository_version}")
+            logger.info(f"[StarflowK8sOperator] Workspace ID: {self.workspace_id}")
+            
+            params = {
+                "workspace_id": self.workspace_id,
+                "url": repository_url,
+                "version": repository_version,
+                'source': source,
+                "page": 0,
+                "page_size": 1,
+            }
+            
+            logger.info("[StarflowK8sOperator] Calling get_piece_repositories_from_workspace_id...")
+            piece_repository_response = self.starflow_client.get_piece_repositories_from_workspace_id(
+                params=params
+            )
+            logger.info(f"[StarflowK8sOperator] get_piece_repositories_from_workspace_id response status: {piece_repository_response.status_code}")
+            
+            piece_repository_data = piece_repository_response.json()
+            logger.info(f"[StarflowK8sOperator] Piece repository data: {piece_repository_data}")
+            
+            # Check if any piece repositories were found
+            if not piece_repository_data.get("data") or len(piece_repository_data["data"]) == 0:
+                logger.warning(f"[StarflowK8sOperator] No piece repositories found for workspace_id={self.workspace_id}, url={repository_url}, version={repository_version}")
+                logger.info("[StarflowK8sOperator] Returning empty secrets dict")
+                return {}
+            
+            piece_repository_id = piece_repository_data["data"][0]["id"]
+            logger.info(f"[StarflowK8sOperator] Piece repository ID: {piece_repository_id}")
+            
+            logger.info("[StarflowK8sOperator] Calling get_piece_secrets...")
+            secrets_response = self.starflow_client.get_piece_secrets(
+                piece_repository_id=piece_repository_id,
+                piece_name=piece_name
+            )
+            logger.info(f"[StarflowK8sOperator] get_piece_secrets response status: {secrets_response.status_code}")
+            
+            if secrets_response.status_code == 404:
+                logger.info("[StarflowK8sOperator] No secrets configured for this piece, returning empty dict")
+                return {}
+            elif secrets_response.status_code != 200:
+                logger.error(f"[StarflowK8sOperator] Error getting piece secrets: {secrets_response.json()}")
+                raise Exception(f"Error getting piece secrets: {secrets_response.json()}")
+
+            secrets_data = secrets_response.json()
+            logger.info(f"[StarflowK8sOperator] Secrets data: {secrets_data}")
+            
+            piece_secrets = {}
+            for e in secrets_data:
+                if not e.get('value') and not e.get('required'):
+                    continue
+                piece_secrets[e.get('name')] = e.get('value')
+            
+            logger.info(f"[StarflowK8sOperator] Piece secrets processed: {list(piece_secrets.keys())}")
+            return piece_secrets
+            
+        except Exception as e:
+            logger.error(f"[StarflowK8sOperator] Error in _get_piece_secrets: {str(e)}")
+            logger.error(f"[StarflowK8sOperator] Exception type: {type(e).__name__}")
+            import traceback
+            logger.error(f"[StarflowK8sOperator] Traceback: {traceback.format_exc()}")
+            raise
 
 
     @staticmethod
@@ -454,30 +529,80 @@ class starflowKubernetesPodOperator(KubernetesPodOperator):
         Runs at the begining of the execute method.
         Pass extra arguments and configuration as environment variables to the pod
         """
-        # Fetch upstream tasks ids and save them in an ENV var
-        upstream_task_ids = [t.task_id for t in self.get_direct_relatives(upstream=True)]
-        self.env_vars.append({
-            'name': 'AIRFLOW_UPSTREAM_TASKS_IDS',
-            'value': str(upstream_task_ids),
-            'value_from': None
-        })
-        # Pass forward the workflow shared storage source name
-        self.env_vars.append({
-            'name': 'STARFLOW_WORKFLOW_SHARED_STORAGE_SOURCE_NAME',
-            'value': str(self.workflow_shared_storage.source.name) if self.workflow_shared_storage else None,
-            'value_from': None
-        })
-        # Save updated piece input kwargs with upstream data to environment variable
-        self.upstream_xcoms_data = self._get_upstream_xcom_data_from_task_ids(task_ids=upstream_task_ids, context=context)
-        self._update_piece_kwargs_with_upstream_xcom()
+        import logging
+        logger = logging.getLogger(__name__)
+        
+        try:
+            logger.info("[StarflowK8sOperator] Fetching upstream task IDs...")
+            # Fetch upstream tasks ids and save them in an ENV var
+            # Check if operator is properly assigned to a DAG
+            try:
+                # Try to access the dag property - this will raise RuntimeError if not assigned
+                dag = self.dag
+                if dag is None:
+                    logger.warning("[StarflowK8sOperator] Operator DAG is None, skipping upstream task IDs")
+                    upstream_task_ids = []
+                else:
+                    upstream_task_ids = [t.task_id for t in self.get_direct_relatives(upstream=True)]
+                    logger.info(f"[StarflowK8sOperator] Upstream task IDs: {upstream_task_ids}")
+            except RuntimeError as e:
+                if "has not been assigned to a DAG yet" in str(e):
+                    logger.warning("[StarflowK8sOperator] Operator not assigned to a DAG, skipping upstream task IDs")
+                    upstream_task_ids = []
+                else:
+                    raise
+            
+            self.env_vars.append({
+                'name': 'AIRFLOW_UPSTREAM_TASKS_IDS',
+                'value': str(upstream_task_ids),
+                'value_from': None
+            })
+            
+            # Pass forward the workflow shared storage source name
+            logger.info("[StarflowK8sOperator] Setting workflow shared storage environment variable...")
+            self.env_vars.append({
+                'name': 'STARFLOW_WORKFLOW_SHARED_STORAGE_SOURCE_NAME',
+                'value': str(self.workflow_shared_storage.source.name) if self.workflow_shared_storage else None,
+                'value_from': None
+            })
+            
+            # Save updated piece input kwargs with upstream data to environment variable
+            logger.info("[StarflowK8sOperator] Getting upstream XCom data...")
+            self.upstream_xcoms_data = self._get_upstream_xcom_data_from_task_ids(task_ids=upstream_task_ids, context=context)
+            logger.info("[StarflowK8sOperator] Upstream XCom data retrieved successfully")
+            
+            logger.info("[StarflowK8sOperator] Updating piece kwargs with upstream XCom...")
+            self._update_piece_kwargs_with_upstream_xcom()
+            logger.info("[StarflowK8sOperator] Piece kwargs updated successfully")
+            
+            logger.info("[StarflowK8sOperator] Getting piece secrets...")
+            piece_secrets = self._get_piece_secrets(
+                repository_url=self.repository_url,
+                repository_version=self.repository_version,
+                piece_name=self.piece_name,
+                source='github'
+            )
+            logger.info(f"[StarflowK8sOperator] Piece secrets retrieved: {list(piece_secrets.keys())}")
+            
+            logger.info("[StarflowK8sOperator] Setting piece secrets environment variables...")
+            for secret_name, secret_value in piece_secrets.items():
+                self.env_vars.append({
+                    'name': secret_name,
+                    'value': str(secret_value),
+                    'value_from': None
+                })
+            logger.info("[StarflowK8sOperator] Piece secrets environment variables set successfully")
+            
+        except Exception as e:
+            logger.error(f"[StarflowK8sOperator] Error in _prepare_execute_environment: {str(e)}")
+            logger.error(f"[StarflowK8sOperator] Exception type: {type(e).__name__}")
+            import traceback
+            logger.error(f"[StarflowK8sOperator] Traceback: {traceback.format_exc()}")
+            raise
+        
+        logger.info("[StarflowK8sOperator] Updating STARFLOW_RUN_PIECE_KWARGS environment variable...")
         self._update_env_var_value_from_name(name='STARFLOW_RUN_PIECE_KWARGS', value=str(self.piece_input_kwargs))
-
-        # # Add pieces secrets to environment variables
-        # piece_secrets = self._get_piece_secrets(
-        #     repository_url=self.repository_url,
-        #     repository_version=self.repository_version,
-        #     piece_name=self.piece_name,
-        # )
+        logger.info("[StarflowK8sOperator] STARFLOW_RUN_PIECE_KWARGS environment variable updated successfully")
         # self.env_vars.append({
         #     "name": "STARFLOW_PIECE_SECRETS",
         #     "value": str(piece_secrets),
@@ -505,19 +630,49 @@ class starflowKubernetesPodOperator(KubernetesPodOperator):
         """
         Code from here onward is executed by the Worker and not by the Scheduler.
         """
-        # TODO change url based on platform configuration
-        self.starflow_client = starflowBackendRestClient(base_url="http://starflow-sa-rest-service.prochain-starflow:8000/")
-        self._prepare_execute_environment(context=context)
-        remote_pod = None
+        import logging
+        logger = logging.getLogger(__name__)
+        
+        # Log at the very beginning to confirm execute method is called
+        logger.info("=" * 80)
+        logger.info(f"[StarflowK8sOperator] EXECUTE METHOD CALLED - task_id: {self.task_id}")
+        logger.info(f"[StarflowK8sOperator] Context type: {type(context)}")
+        logger.info(f"[StarflowK8sOperator] Context keys: {list(context.keys()) if hasattr(context, 'keys') else 'No keys method'}")
+        logger.info("=" * 80)
+        
         try:
+            logger.info(f"[StarflowK8sOperator] Starting execution for task_id: {self.task_id}")
+            
+            # Initialize remote_pod early to avoid UnboundLocalError in finally block
+            remote_pod = None
+            
+            # TODO change url based on platform configuration
+            logger.info("[StarflowK8sOperator] Initializing Starflow backend client...")
+            self.starflow_client = starflowBackendRestClient(base_url="http://starflow-rest-service.starflow:8000/")
+            logger.info("[StarflowK8sOperator] Starflow backend client initialized successfully")
+            
+            logger.info("[StarflowK8sOperator] Preparing execute environment...")
+            self._prepare_execute_environment(context=context)
+            logger.info("[StarflowK8sOperator] Execute environment prepared successfully")
+            logger.info("[StarflowK8sOperator] Building pod request object...")
             self.pod_request_obj = self.build_pod_request_obj(context)
+            logger.info("[StarflowK8sOperator] Pod request object built successfully")
+            
+            logger.info("[StarflowK8sOperator] Creating/getting pod...")
             self.pod = self.get_or_create_pod(  # must set `self.pod` for `on_kill`
                 pod_request_obj=self.pod_request_obj,
                 context=context,
             )
+            logger.info(f"[StarflowK8sOperator] Pod created successfully: {self.pod.metadata.name}")
+            
             # get remote pod for use in cleanup methods
+            logger.info("[StarflowK8sOperator] Finding remote pod...")
             remote_pod = self.find_pod(self.pod.metadata.namespace, context=context)
+            logger.info("[StarflowK8sOperator] Remote pod found successfully")
+            
+            logger.info("[StarflowK8sOperator] Awaiting pod start...")
             self.await_pod_start(pod=self.pod)
+            logger.info("[StarflowK8sOperator] Pod started successfully")
 
             if self.get_logs:
                 self.pod_manager.fetch_container_logs(

{starflow_py-0.49.0 → starflow_py-0.56.0}/src/starflow/task.py

@@ -66,10 +66,26 @@ class Task(object):
         self.container_resources = container_resources
 
         # Get deploy mode
-        self.deploy_mode = os.environ.get('STARFLOW_DEPLOY_MODE')
+        self.deploy_mode = os.environ.get('STARFLOW_DEPLOY_MODE', 'prod')
 
         # Set up task operator
         self._task_operator = self._set_operator()
+    
+    def _construct_image_url(self) -> str:
+        """Construct the Docker image URL from repository_url and piece name"""
+        # Extract the base image name from repository_url
+        # repository_url is like: ghcr.io/prochain-star-atlas/prochain-starflow-pieces:0.1.75-group0
+        # We need to construct: ghcr.io/prochain-star-atlas/prochain-starflow-pieces:0.1.75-group0
+        
+        # For now, use the repository_url as the base and append the piece name
+        # This matches the pattern we see in the logs
+        base_image = self.repository_url
+        if ':' in base_image:
+            # If it already has a tag, use it as is
+            return base_image
+        else:
+            # If no tag, append the version
+            return f"{base_image}:{self.repository_version}"
 
     def _set_operator(self) -> BaseOperator:
         """
@@ -96,7 +112,7 @@ class Task(object):
             container_resources=self.container_resources,
             # ----------------- Kubernetes -----------------
             namespace='airflow',
-            image=self.piece.get("source_image"),
+            image=self.piece.get("source_image") or self._construct_image_url(),
             image_pull_policy='IfNotPresent',
             name=f"airflow-worker-pod-{self.task_id}",
             startup_timeout_seconds=600,

{starflow_py-0.49.0 → starflow_py-0.56.0}/src/starflow/task_branch.py

@@ -70,7 +70,7 @@ class TaskBranch(object):
         self.container_resources = container_resources
 
         # Get deploy mode
-        self.deploy_mode = os.environ.get('STARFLOW_DEPLOY_MODE')
+        self.deploy_mode = os.environ.get('STARFLOW_DEPLOY_MODE', 'prod')
 
         # Set up task operator
         self._task_operator = self._set_operator()
@@ -104,7 +104,7 @@ class TaskBranch(object):
             namespace='airflow',
             image=self.piece.get("source_image"),
             image_pull_policy='IfNotPresent',
-            name=f"airflow-worker-pod-{self.task_id}",
+            name=f"airflow-worker-pod-{self.task_id.lower().replace('_', '-')}",
             startup_timeout_seconds=600,
             annotations={"sidecar.istio.io/inject": "false"}, # TODO - remove this when istio is working with airflow k8s pod
             # cmds=["/bin/bash"],

{starflow_py-0.49.0 → starflow_py-0.56.0}/src/starflow_py.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: starflow-py
-Version: 0.49.0
+Version: 0.56.0
 Summary: Python package for starflow.
 Author-email: Alexandre Lazar <alexandre.lazar@gmail.com>
 License: 

starflow_py-0.49.0/src/starflow/VERSION

@@ -1 +0,0 @@
-0.49.0