7373distribution_name = generate_random_name (prefix = "llama-stack-distribution" )
7474
7575
76- def _cleanup_s3_files (
77- bucket_name : str ,
78- endpoint_url : str ,
79- region : str ,
80- access_key_id : str ,
81- secret_access_key : str ,
82- ) -> None :
83- """
84- Clean up files from S3 bucket that were uploaded during tests.
85-
86- Args:
87- bucket_name: S3 bucket name
88- endpoint_url: S3 endpoint URL
89- region: S3 region
90- access_key_id: AWS access key ID
91- secret_access_key: AWS secret access key
92- """
93-
94- try :
95- import boto3
96- from botocore .exceptions import ClientError
97-
98- s3_client = boto3 .client (
99- service_name = "s3" ,
100- endpoint_url = endpoint_url ,
101- aws_access_key_id = access_key_id ,
102- aws_secret_access_key = secret_access_key ,
103- region_name = region ,
104- )
105-
106- response = s3_client .list_objects_v2 (Bucket = bucket_name )
107-
108- if "Contents" not in response :
109- LOGGER .info ("No files found to clean up from S3" )
110- return
111-
112- # We only want to delete files that start with "file-"
113- for obj in response ["Contents" ]:
114- key = obj ["Key" ]
115- if key .startswith ("file-" ):
116- s3_client .delete_object (Bucket = bucket_name , Key = key )
117- LOGGER .debug (f"Deleted file from S3: { key } " )
118-
119- response = s3_client .list_objects_v2 (Bucket = bucket_name )
120-
121- if "Contents" not in response :
122- LOGGER .info ("No files found to clean up from S3" )
123- return
124-
125- except ClientError as e :
126- LOGGER .warning (f"Failed to clean up S3 files: { e } " )
127-
128-
12976@pytest .fixture (scope = "class" )
13077def enabled_llama_stack_operator (dsc_resource : DataScienceCluster ) -> Generator [DataScienceCluster , Any , Any ]:
13178 with update_components_in_dsc (
@@ -385,11 +332,6 @@ def unprivileged_llama_stack_distribution(
385332 enabled_llama_stack_operator : DataScienceCluster ,
386333 request : FixtureRequest ,
387334 llama_stack_server_config : dict [str , Any ],
388- ci_s3_bucket_name : str ,
389- ci_s3_bucket_endpoint : str ,
390- ci_s3_bucket_region : str ,
391- aws_access_key_id : str ,
392- aws_secret_access_key : str ,
393335 unprivileged_llama_stack_distribution_secret : Secret ,
394336 unprivileged_postgres_deployment : Deployment ,
395337 unprivileged_postgres_service : Service ,
@@ -406,25 +348,6 @@ def unprivileged_llama_stack_distribution(
406348 lls_dist .wait_for_status (status = LlamaStackDistribution .Status .READY , timeout = 600 )
407349 yield lls_dist
408350
409- try :
410- env_vars = llama_stack_server_config .get ("containerSpec" , {}).get ("env" , [])
411- enable_s3 = any (env .get ("name" ) == "ENABLE_S3" and env .get ("value" ) == "s3" for env in env_vars )
412-
413- if enable_s3 :
414- try :
415- _cleanup_s3_files (
416- bucket_name = ci_s3_bucket_name ,
417- endpoint_url = ci_s3_bucket_endpoint ,
418- region = ci_s3_bucket_region ,
419- access_key_id = aws_access_key_id ,
420- secret_access_key = aws_secret_access_key ,
421- )
422- except Exception as e : # noqa: BLE001
423- LOGGER .warning (f"Failed to clean up S3 files: { e } " )
424-
425- except Exception as e : # noqa: BLE001
426- LOGGER .warning (f"Failed to clean up S3 files: { e } " )
427-
428351
429352@pytest .fixture (scope = "class" )
430353def llama_stack_distribution (
@@ -433,11 +356,6 @@ def llama_stack_distribution(
433356 enabled_llama_stack_operator : DataScienceCluster ,
434357 request : FixtureRequest ,
435358 llama_stack_server_config : dict [str , Any ],
436- ci_s3_bucket_name : str ,
437- ci_s3_bucket_endpoint : str ,
438- ci_s3_bucket_region : str ,
439- aws_access_key_id : str ,
440- aws_secret_access_key : str ,
441359 llama_stack_distribution_secret : Secret ,
442360 postgres_deployment : Deployment ,
443361 postgres_service : Service ,
@@ -453,25 +371,6 @@ def llama_stack_distribution(
453371 lls_dist .wait_for_status (status = LlamaStackDistribution .Status .READY , timeout = 600 )
454372 yield lls_dist
455373
456- try :
457- env_vars = llama_stack_server_config .get ("containerSpec" , {}).get ("env" , [])
458- enable_s3 = any (env .get ("name" ) == "ENABLE_S3" and env .get ("value" ) == "s3" for env in env_vars )
459-
460- if enable_s3 :
461- try :
462- _cleanup_s3_files (
463- bucket_name = ci_s3_bucket_name ,
464- endpoint_url = ci_s3_bucket_endpoint ,
465- region = ci_s3_bucket_region ,
466- access_key_id = aws_access_key_id ,
467- secret_access_key = aws_secret_access_key ,
468- )
469- except Exception as e : # noqa: BLE001
470- LOGGER .warning (f"Failed to clean up S3 files: { e } " )
471-
472- except Exception as e : # noqa: BLE001
473- LOGGER .warning (f"Failed to clean up S3 files: { e } " )
474-
475374
476375def _get_llama_stack_distribution_deployment (
477376 client : DynamicClient ,
@@ -642,11 +541,37 @@ def _create_llama_stack_client(
642541 http_client = http_client ,
643542 )
644543 wait_for_llama_stack_client_ready (client = client )
544+ existing_file_ids = {f .id for f in client .files .list ().data }
545+
645546 yield client
547+
548+ _cleanup_files (client = client , existing_file_ids = existing_file_ids )
646549 finally :
647550 http_client .close ()
648551
649552
553+ def _cleanup_files (client : LlamaStackClient , existing_file_ids : set [str ]) -> None :
554+ """Delete files created during test execution via the LlamaStack files API.
555+
556+ Only deletes files whose IDs were not present before the test ran,
557+ avoiding interference with other test sessions.
558+
559+ Args:
560+ client: The LlamaStackClient used during the test
561+ existing_file_ids: File IDs that existed before the test started
562+ """
563+ try :
564+ for file in client .files .list ().data :
565+ if file .id not in existing_file_ids :
566+ try :
567+ client .files .delete (file_id = file .id )
568+ LOGGER .debug (f"Deleted file: { file .id } " )
569+ except Exception as e :
570+ LOGGER .warning (f"Failed to delete file { file .id } : { e } " )
571+ except Exception as e :
572+ LOGGER .warning (f"Failed to clean up files: { e } " )
573+
574+
650575@pytest .fixture (scope = "class" )
651576def unprivileged_llama_stack_client (
652577 unprivileged_llama_stack_test_route : Route ,
0 commit comments